kvesteri / postgresql-audit

Audit trigger for PostgreSQL
BSD 2-Clause "Simplified" License
126 stars 28 forks source link

Refactor Audit Table query builder for manual Alembic migration support #108

Closed mparent61 closed 8 months ago

mparent61 commented 8 months ago

This is a small refactor of the "audit table" query logic, to allow an easy workaround for Alembic migrations.

I've seen many issues/PRs already about a lack of Alembic support, but since they appear to be both complicated and stuck in review, I needed a simpler approach in order to go to production with this library.

This just allows the core audit table query logic to easily be shared by custom Alembic helper functions (not in this PR, but provided as a workaround / example for other developers struggling with Alembic integration):

def enable_or_update_model_auditing(op, model):
    """Run PGAudit stored function to enable auditing specific Model.

    This function is designed to be run whenever auditing is enabled or updated (ex: excluded column list is modified).
    """
    from my_auditing import versioning_manager

    assert hasattr(
        model, "__versioned__"
    ), f"Auditing requires `__versioned__` class variable on model `{model}`"

    op.execute(
        versioning_manager.build_audit_table_query(
            model.__table__, model.__versioned__.get("exclude")
        )
    )

def disable_model_auditing(op, model):
    """Disable table-specific auditing machinery. The inverse of `enable_or_update_model_auditing`."""
    tablename = model.__tablename__

    op.execute(f"DROP TRIGGER IF EXISTS audit_trigger_insert ON {tablename}")
    op.execute(f"DROP TRIGGER IF EXISTS audit_trigger_update ON {tablename}")
    op.execute(f"DROP TRIGGER IF EXISTS audit_trigger_delete ON {tablename}")
    op.execute(f"DROP TRIGGER IF EXISTS audit_trigger_row ON {tablename}")

Then, I just call either one in my Alembic migration scripts, whenever I add/update/remove auditing configurations:

def upgrade():
    from my_models import User

    enable_or_update_model_auditing(op, User)

def downgrade():
    from my_models import User

    disable_model_auditing(op, User)