feat: integrate table comment extraction for models

Generated-by: aiautocommit

activemodel/base_model.py

@@ -41,6 +41,8 @@ class BaseModel(SQLModel):
 
     # TODO implement actually calling these hooks
 
+    __table_args__ = None
+
     @classmethod
     def __init_subclass__(cls, **kwargs):
         "Setup automatic sqlalchemy lifecycle events for the class"
@@ -53,6 +55,8 @@ def __init_subclass__(cls, **kwargs):
         # sa args to it it persisted to the sql table comments
         set_config_value(model=cls, parameter="use_attribute_docstrings", value=True)
 
+        cls._apply_class_doc()
+
         def event_wrapper(method_name: str):
             """
             This does smart heavy lifting for us to make sqlalchemy lifecycle events nicer to work with:
@@ -102,8 +106,24 @@ def wrapper(mapper: Mapper, connection: Connection, target: BaseModel):
         event.listen(cls, "after_insert", event_wrapper("after_save"))
         event.listen(cls, "after_update", event_wrapper("after_save"))
 
-    # def foreign_key()
-    # table.id
+        # def foreign_key()
+        # table.id
+
+    @classmethod
+    def _apply_class_doc(cls):
+        "Pull the class-level docstring into a table comment"
+
+        doc = cls.__doc__.strip() if cls.__doc__ else None
+
+        if doc:
+            table_args = getattr(cls, "__table_args__", None)
+
+            if table_args is None:
+                cls.__table_args__ = {"comment": doc}
+            elif isinstance(table_args, dict):
+                table_args.setdefault("comment", doc)
+            else:
+                raise ValueError("Unexpected __table_args__ type")
 
     # TODO no type check decorator here
     @declared_attr

test/comments_test.py

@@ -0,0 +1,47 @@
+"""
+Test database comments integration
+"""
+
+from activemodel import BaseModel
+from activemodel.mixins.timestamps import TimestampsMixin
+from activemodel.mixins.typeid import TypeIDMixin
+from activemodel.session_manager import get_session
+from sqlmodel import text
+
+
+class ExampleWithoutComments(
+    BaseModel, TimestampsMixin, TypeIDMixin("ex2"), table=True
+):
+    pass
+
+
+TABLE_COMMENT = "Expected table comment"
+
+
+class ExampleWithComments(BaseModel, TimestampsMixin, TypeIDMixin("ex"), table=True):
+    """Expected table comment"""
+
+
+def test_comment(create_and_wipe_database):
+    assert ExampleWithComments.__doc__
+    assert not ExampleWithoutComments.__doc__
+
+    connection = create_and_wipe_database
+    with get_session() as session:
+        result = session.execute(
+            text("""
+            SELECT obj_description('example_with_comments'::regclass, 'pg_class') AS table_comment;
+        """)
+        )
+        table_comment = result.fetchone()[0]
+        assert table_comment == "Expected table comment"
+
+        # session.exec(
+        #     text("""
+        #     SELECT col_description('example_with_comments'::regclass, ordinal_position) AS column_comment
+        #     FROM information_schema.columns
+        #     WHERE table_name = 'example_with_comments' AND column_name = 'your_column';
+        # """)
+        # )
+        # column_comment = cursor.fetchone()[0]
+        # assert column_comment == "Expected column comment"

test/conftest.py

@@ -1,5 +1,19 @@
+import pytest
+
 import activemodel
+from activemodel.session_manager import get_engine
+from sqlmodel import SQLModel
 
-from .utils import database_url
+from .utils import database_url, temporary_tables
 
 activemodel.init(database_url())
+
+SQLModel.metadata.drop_all(
+    bind=get_engine(),
+)
+
+
+@pytest.fixture(scope="function")
+def create_and_wipe_database():
+    with temporary_tables():
+        yield

test/orm_test.py

@@ -2,26 +2,11 @@
 Test core ORM functions
 """
 
-from contextlib import contextmanager
+from test.utils import temporary_tables
 
-from activemodel import BaseModel, get_engine
+from activemodel import BaseModel
 from activemodel.mixins.timestamps import TimestampsMixin
 from activemodel.mixins.typeid import TypeIDMixin
-from sqlmodel import SQLModel
-
-
-@contextmanager
-def temporary_tables():
-    SQLModel.metadata.create_all(get_engine())
-
-    try:
-        yield
-    finally:
-        SQLModel.metadata.drop_all(
-            # tables=[SQLModel.metadata.tables[AIVisitNote.__tablename__]],
-            bind=get_engine(),
-        )
-
 
 EXAMPLE_TABLE_PREFIX = "test_record"
 

test/utils.py

@@ -1,4 +1,8 @@
 import os
+from contextlib import contextmanager
+
+from activemodel import get_engine
+from sqlmodel import SQLModel
 
 
 def database_url():
@@ -14,3 +18,15 @@ def database_url():
     # https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls
     # without this, psycopg2 would be used, which is not intended!
     return url.replace("postgresql://", "postgresql+psycopg://")
+
+
+@contextmanager
+def temporary_tables():
+    SQLModel.metadata.create_all(get_engine())
+
+    try:
+        yield
+    finally:
+        SQLModel.metadata.drop_all(
+            bind=get_engine(),
+        )