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

Author: DanCardin

docs/Makefile

@@ -3,7 +3,7 @@
 
 # You can set these variables from the command line, and also
 # from the environment for the first two.
-SPHINXOPTS    ?=-W -a -E
+SPHINXOPTS    ?= -a -E
 SPHINXBUILD   ?= sphinx-build
 SOURCEDIR     = source
 BUILDDIR      = build

docs/source/api.rst

@@ -12,3 +12,17 @@ API
 
 .. autofunction:: sqlalchemy_postgresql_audit.uninstall_audit_triggers
     :noindex:
+
+Declarative API
+===============
+An alternative API for the use and enablement of auditing functionality
+can be used more directly on existing models/tables.
+
+.. autofunction:: sqlalchemy_postgresql_audit.audit_model
+    :noindex:
+
+.. autofunction:: sqlalchemy_postgresql_audit.create_audit_model
+    :noindex:
+
+.. autofunction:: sqlalchemy_postgresql_audit.create_audit_table
+    :noindex:

setup.cfg

@@ -7,3 +7,15 @@ max-line-length = 100
 
 [black]
 line-length = 100
+
+[isort]
+profile = black
+known_first_party = sqlalchemy_postgresql_audit,tests
+line_length = 100
+float_to_top=true
+order_by_type = false
+use_parentheses = true
+
+[tool:pytest]
+doctest_optionflags = NORMALIZE_WHITESPACE IGNORE_EXCEPTION_DETAIL ELLIPSIS
+addopts = --doctest-modules --ff --strict-markers

src/sqlalchemy_postgresql_audit/__init__.py

@@ -5,8 +5,12 @@
     "enable",
     "install_audit_triggers",
     "uninstall_audit_triggers",
+    "audit_model",
+    "create_audit_model",
+    "create_audit_table",
 ]
 
-from .session import set_session_vars
-from .plugin import enable
+from .declarative import audit_model, create_audit_model, create_audit_table
 from .ddl import install_audit_triggers, uninstall_audit_triggers
+from .plugin import enable
+from .session import set_session_vars

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/declarative.py

@@ -0,0 +1,155 @@
+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 as create_raw_audit_table
+
+
+def default_primary_key():
+    return Column(
+        'audit_uuid',
+        UUID(as_uuid=True),
+        primary_key=True,
+        default=uuid.uuid4,
+        server_default=text("uuid_generate_v4()"),
+    )
+
+
+def audit_model(_func=None, *, enabled=True, primary_key=default_primary_key, **spec):
+    """Decorate a model to automatically enable audit modeling.
+
+    Arguments:
+        enabled: Defaults to true, enables auditing.
+        primary_key: Default to a uuid primary key. Can be disabled by using `None`.
+
+    By default, automatically enables the auditing in addition to hooking
+    up the actual audit machinery.
+
+    Additionally, leaves a reference to the audit model's own sqlachemy model
+    on the ``__audit_cls__`` attribute of the decorated class.
+
+    Examples:
+        >>> from sqlalchemy import Column, types
+        >>> from sqlalchemy.ext.declarative import declarative_base
+        >>> from sqlalchemy_postgresql_audit import audit_model
+
+        >>> Base = declarative_base()
+
+        >>> @audit_model
+        ... class Foo(Base):
+        ...     __tablename__ = 'foo'
+        ...     id = Column(types.Integer(), primary_key=True)
+
+        >>> Foo.__audit_cls__
+        <class '...FooAudit'>
+
+        >>> @audit_model(enabled=False)
+        ... class Bar(Base):
+        ...     __tablename__ = 'bar'
+        ...     id = Column(types.Integer(), primary_key=True)
+    """
+    def decorated(model_cls):
+        model = create_audit_model(model_cls, enabled=enabled, **spec)
+        if model:
+            model_cls.__audit_cls__ = model
+
+        return model_cls
+
+    if _func is None:
+        return decorated
+    return decorated(_func)
+
+
+def create_audit_model(model_cls, *, enabled=True, primary_key=default_primary_key, **spec):
+    """Create an SQLAlchemy declarative Model class for the given `model_cls`.
+
+    Arguments:
+        model_cls: The SQLAlchemy model being audited
+        enabled: Defaults to true, enables auditing.
+        primary_key: Default to a uuid primary key. Can be disabled by using `None`.
+
+    Examples:
+        >>> from sqlalchemy import Column, types
+        >>> from sqlalchemy.ext.declarative import declarative_base
+        >>> from sqlalchemy_postgresql_audit import create_audit_model
+
+        >>> Base = declarative_base()
+
+        >>> class Foo(Base):
+        ...     __tablename__ = 'foo'
+        ...     id = Column(types.Integer(), primary_key=True)
+
+        >>> class Bar(Base):
+        ...     __tablename__ = 'bar'
+        ...     id = Column(types.Integer(), primary_key=True)
+
+        >>> class Baz(Base):
+        ...     __tablename__ = 'baz'
+        ...     id = Column(types.Integer(), primary_key=True)
+
+        >>> AuditModel = create_audit_model(Foo)
+        >>> AuditModel3 = create_audit_model(Baz, primary_key=default_primary_key)
+        >>> create_audit_model(Bar, enabled=False)
+    """
+    base_table = model_cls.__table__
+    metadata = model_cls.metadata
+
+    table = create_audit_table(base_table, metadata, enabled=enabled, primary_key=primary_key, **spec)
+    if table is None:
+        return
+
+    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=default_primary_key, **spec):
+    """Create an audit SQLAlchemy ``Table`` for a given `Table` instance.
+
+    Arguments:
+        table: The SQLAlchemy `Table` to audit.
+        metadata: The `SQLAlchemy` metadata on which to attach the table.
+        enabled: Defaults to true, enables auditing.
+        primary_key: Default to a uuid primary key. Can be disabled by using `None`.
+        spec: Optional auditing spec options.
+
+    Examples:
+        >>> from sqlalchemy import MetaData, Table
+        >>> from sqlalchemy_postgresql_audit import create_audit_table
+
+        >>> meta = MetaData()
+
+        >>> foo_table = Table('foo', meta)
+        >>> audit_table1 = create_audit_table(foo_table, meta)
+
+        >>> baz_table = Table('baz', meta)
+        >>> audit_table3 = create_audit_table(baz_table, meta, primary_key=None)
+
+        >>> bar_table = Table('bar', meta)
+        >>> create_audit_table(bar_table, meta, enabled=False)
+    """
+    existing_info = table.info
+    existing_info['audit.options'] = {'enabled': enabled, **spec}
+    
+    primary_key_column = primary_key() if primary_key else None
+
+    return create_raw_audit_table(table, metadata, primary_key=primary_key_column)
+
+
+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/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_declarative.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

tox.ini

@@ -13,7 +13,7 @@ basepython =
 
 commands =
     pip install sqlalchemy-postgresql-audit[testing]
-    pytest --cov sqlalchemy_postgresql_audit --cov-report= -v tests {posargs:}
+    pytest --cov sqlalchemy_postgresql_audit --cov-report= -v src tests {posargs:}
 
 setenv =
     COVERAGE_FILE=tmp/.coverage.{envname}