Merge pull request #578 from TotallyNotRobots/support-sql-orm

Add support for sqlalchemy ORM objects

Author: linuxdaemon

cloudbot/plugin.py

@@ -6,7 +6,7 @@
 from functools import partial
 from operator import attrgetter
 from pathlib import Path
-from typing import Dict, List, MutableMapping, Optional
+from typing import Dict, List, MutableMapping, Optional, Type, cast
 from weakref import WeakValueDictionary
 
 import sqlalchemy
@@ -54,6 +54,9 @@ def find_tables(code):
         ):
             # if it's a Table, and it's using our metadata, append it to the list
             tables.append(obj)
+        elif isinstance(obj, type) and issubclass(obj, database.Base):
+            obj = cast(Type[database.Base], obj)
+            tables.append(obj.__table__)
 
     return tables
 

tests/core_tests/test_plugin_manager.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+import sqlalchemy as sa
 from sqlalchemy import Column, String, Table, inspect
 
 from cloudbot import hook
@@ -56,6 +57,28 @@ def test_get_plugin(mock_manager):
     assert mock_manager.find_plugin("test") is None
 
 
+def test_find_tables(mock_manager):
+    file_path = mock_manager.bot.plugin_dir / "test.py"
+    file_name = file_path.name
+
+    class TestTable(database.Base):
+        __tablename__ = "foo"
+
+        test_id = sa.Column(sa.Integer, primary_key=True)
+
+    mod = MockModule(tbl=TestTable)
+    obj = Plugin(
+        str(file_path),
+        file_name,
+        "test",
+        mod,
+    )
+
+    table_names = [t.name for t in obj.tables]
+
+    assert table_names == ["foo"]
+
+
 def test_can_load(mock_manager):
     mock_manager.bot.config.clear()