feat: Add a more ergonomic API for enabling auditing on models.

Author: DanCardin

src/sqlalchemy_postgresql_audit/__init__.py

@@ -7,6 +7,7 @@
     "uninstall_audit_triggers",
 ]
 
-from .session import set_session_vars
-from .plugin import enable
+from .api import audit_model
 from .ddl import install_audit_triggers, uninstall_audit_triggers
+from .plugin import enable
+from .session import set_session_vars

src/sqlalchemy_postgresql_audit/api.py

@@ -0,0 +1,62 @@
+import functools
+import uuid
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy import Column, text
+
+try:
+    from sqlalchemy.orm.decl_api import DeclarativeMeta
+except ImportError:
+    from sqlalchemy.ext.declarative.api import DeclarativeMeta
+
+from sqlalchemy_postgresql_audit.event_listeners.sqlalchemy import create_audit_table
+
+
+def audit_model(_func=None, *, enabled=True, **spec):
+    def decorated(model_cls):
+        create_audit_model(model_cls, enabled=enabled, **spec)
+
+        return model_cls
+
+    if _func is None:
+        return decorated
+    return decorated(_func)
+
+
+def create_audit_model(model_cls, *, enabled=True, primary_key=None, **spec):
+    table = model_cls.__table__
+    metadata = model_cls.metadata
+
+    table = create_audit_table(table, metadata, enabled=enabled, primary_key=primary_key, **spec)
+    model_base = _find_model_base(model_cls)
+
+    cls = type(
+        '{model_cls}Audit'.format(model_cls=model_cls.__name__),
+        (model_base,),
+        {'__table__': table},
+    )
+
+    return cls
+
+
+def create_audit_table(table, metadata, *, enabled=True, primary_key=None, **spec):
+    existing_info = table.info
+    existing_info['audit.options'] = {'enabled': enabled, **spec}
+
+    if primary_key is None:
+        primary_key = Column(
+            'audit_uuid',
+            UUID(as_uuid=True),
+            primary_key=True,
+            default=uuid.uuid4,
+            server_default=text("uuid_generate_v4()"),
+        )
+
+    return create_audit_table(table, metadata, primary_key=primary_key)
+
+
+def _find_model_base(model_cls):
+    for cls in model_cls.__mro__:
+        if isinstance(cls, DeclarativeMeta) and not hasattr(cls, '__mapper__'):
+            return cls
+
+    raise ValueError("Invalid model, does not subclass a `DeclarativeMeta`.")

src/sqlalchemy_postgresql_audit/ddl.py

@@ -56,6 +56,7 @@ def get_create_trigger_ddl(
         # or it is one of our session settings values
         else:
             if col.name in (
+                "audit_uuid",
                 "audit_operation",
                 "audit_operation_timestamp",
                 "audit_current_user",

src/sqlalchemy_postgresql_audit/event_listeners/__init__.py

@@ -1,7 +1,6 @@
 import threading
 
-from sqlalchemy import Table
-from sqlalchemy.events import event
+from sqlalchemy import Table, event
 
 _event_listeners_enabled = False
 
@@ -17,20 +16,19 @@ def enable_event_listeners():
 
 
 def _enable_sqlalchemy_event_listeners():
-    from sqlalchemy_postgresql_audit.event_listeners.sqlalchemy import (
-        create_audit_table,
-    )
+    from sqlalchemy_postgresql_audit.event_listeners.sqlalchemy import \
+        create_audit_table
 
     event.listens_for(Table, "after_parent_attach")(create_audit_table)
 
 
 def _enable_alembic_event_listeners():
     try:
-        from sqlalchemy_postgresql_audit.event_listeners.alembic import (
-            compare_for_table,
-        )
         from alembic.autogenerate.compare import comparators
 
+        from sqlalchemy_postgresql_audit.event_listeners.alembic import \
+            compare_for_table
+
         comparators.dispatch_for("table")(compare_for_table)
     except ImportError:
         pass

src/sqlalchemy_postgresql_audit/event_listeners/sqlalchemy.py

@@ -18,7 +18,7 @@
 )
 
 
-def create_audit_table(target, parent):
+def create_audit_table(target, parent, primary_key=None):
     """Create an audit table and generate procedure/trigger DDL.
 
     Naming conventions can be defined for a few of the named elements:
@@ -83,21 +83,29 @@ def create_audit_table(target, parent):
         "schema": audit_spec["schema"] or "public",
     }
 
-    columns = [
-        Column(col.name, col.type, nullable=True) for col in target.columns.values()
-    ]
+    column_elements = []
+    if primary_key is not None:
+        column_elements.append(primary_key)
+
+    column_elements.extend([
+        Column("audit_operation", String(1), nullable=False),
+        Column("audit_operation_timestamp", DateTime, nullable=False),
+        Column("audit_current_user", String(64), nullable=False),
+    ])
+
     session_setting_columns = [col.copy() for col in audit_spec["session_settings"]]
     for col in session_setting_columns:
         col.name = "audit_{}".format(col.name)
+    column_elements.extend(session_setting_columns)
 
-    column_elements = session_setting_columns + columns
+    table_columns = [
+        Column(col.name, col.type, nullable=True) for col in target.columns.values()
+    ]
+    column_elements.extend(table_columns)
 
     audit_table = Table(
         audit_table_name,
         target.metadata,
-        Column("audit_operation", String(1), nullable=False),
-        Column("audit_operation_timestamp", DateTime, nullable=False),
-        Column("audit_current_user", String(64), nullable=False),
         *column_elements,
         schema=audit_spec["schema"]
     )
@@ -119,3 +127,4 @@ def create_audit_table(target, parent):
 
     audit_table.info["audit.is_audit_table"] = True
     target.info["audit.is_audited"] = True
+    return audit_table

tests/unit_tests/test_audit_model.py

@@ -0,0 +1,44 @@
+from sqlalchemy import Column, Integer, MetaData, Table
+from sqlalchemy.ext.declarative import declarative_base
+
+import sqlalchemy_postgresql_audit
+
+
+def setup():
+    sqlalchemy_postgresql_audit.enable()
+
+
+def test_vanilla_model():
+    Base = declarative_base()
+    metadata = Base.metadata
+
+    @sqlalchemy_postgresql_audit.audit_model
+    class Model(Base):
+        __tablename__ = 'foo'
+
+        id = Column("id", Integer, primary_key=True)
+
+    audit_table = metadata.tables["foo_audit"]
+
+    assert audit_table.info["audit.is_audit_table"]
+    assert Model.__table__.info["audit.is_audited"]
+
+
+def test_model_with_info():
+    Base = declarative_base()
+    metadata = Base.metadata
+
+    @sqlalchemy_postgresql_audit.audit_model
+    class Model(Base):
+        __tablename__ = 'foo'
+        __table_args__ = {
+                'info': {'example': 4}
+            }
+
+        id = Column("id", Integer, primary_key=True)
+
+    audit_table = metadata.tables["foo_audit"]
+
+    assert audit_table.info["audit.is_audit_table"]
+    assert Model.__table__.info["audit.is_audited"]
+    assert 'example' in  Model.__table__.info