flask-postgresql-audit 1.0.7__tar.gz → 1.0.8__tar.gz

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flask-postgresql_audit
-Version: 1.0.7
+Version: 1.0.8
 Summary: Versioning and auditing extension for PostgreSQL and SQLAlchemy.
 Project-URL: Homepage, https://github.com/higamigu/flask-postgresql-audit
 Author-email: Higamigu <higamigu@gmail.com>, Konsta Vesterinen <konsta@fastmonkeys.com>

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/flask_postgresql_audit/__init__.py

@@ -1,4 +1,4 @@
 # ruff: noqa: F401
 from .core import Audit, PostgreSQLAudit
 
-__version__ = "1.0.7"
+__version__ = "1.0.8"

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/flask_postgresql_audit/alembic/entities.py

@@ -55,6 +55,15 @@ def get_pk_values(**context):
     )
 
 
+def get_table_relid(**context):
+    sql = parse_template("get_table_relid.sql", **context)
+    return PGFunction(
+        schema=context["schema_name"],
+        signature="get_table_relid(tablename text)",
+        definition=sql["definition"],
+    )
+
+
 def get_setting_factory(**context):
     sql = parse_template("get_setting.sql", **context)
     return PGFunction(

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/flask_postgresql_audit/alembic/setup.py

@@ -63,6 +63,7 @@ def setup_db(audit: "PostgreSQLAudit"):
 def register_core_entities(audit: "PostgreSQLAudit"):
     audit.pg_audit_entities.add(entities.btree_gist)
     audit.pg_audit_entities.add(entities.get_pk_values(**audit.context))
+    audit.pg_audit_entities.add(entities.get_table_relid(**audit.context))
     audit.pg_audit_entities.add(entities.get_setting_factory(**audit.context))
     audit.pg_audit_entities.add(entities.jsonb_subtract_factory(**audit.context))
     audit.pg_audit_entities.add(entities.create_activity_factory(**audit.context))

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/flask_postgresql_audit/core.py

@@ -27,6 +27,12 @@ class ImproperlyConfigured(Exception):
 
 class Audit:
     __audit_args__: dict = {}
+    __audit_schema__: str = ""
+
+    __audit_relid__: orm.Mapped[int]
+    __audit_rowkey__: orm.Mapped[list[str]]
+
+    __audit_configured__: t.ClassVar[bool]
 
     if t.TYPE_CHECKING:
         __table__: sa.Table
@@ -34,10 +40,33 @@ class Audit:
 
         def __init__(self, **kw: t.Any): ...
 
-    def __init_subclass__(cls) -> None:
+    def __init_subclass__(cls, *args, **kwargs) -> None:
         if not hasattr(cls, "__audit_args__"):
             cls.__audit_args__ = {}
 
+        event.listen(orm.Mapper, "mapper_configured", cls.__init_audit__)
+        super().__init_subclass__(*args, **kwargs)
+
+    @classmethod
+    def __init_audit__(cls, mapper: orm.Mapper[t.Self], class_: type[t.Self]):
+        is_configured = getattr(class_, "__audit_configured__", False)
+        has_schema = hasattr(class_, "__audit_schema__")
+        if not is_configured and has_schema:
+            fn_ = sa.func
+            if class_.__audit_schema__ and class_.__audit_schema__ != "public":
+                fn_ = getattr(sa.func, class_.__audit_schema__)
+
+            table = class_.__table__
+            jsonb = sa.func.to_jsonb(table.table_valued())
+
+            relid = orm.column_property(fn_.get_table_relid(class_.__tablename__))
+            rowkey = orm.column_property(fn_.get_pk_values(relid, jsonb))
+
+            mapper.add_property("__audit_relid__", relid)
+            mapper.add_property("__audit_rowkey__", rowkey)
+
+            class_.__audit_configured__ = True
+
 
 def get_modified_columns(obj: orm.DeclarativeBase):
     modified: set[sa.ColumnElement[t.Any]] = set()
@@ -65,6 +94,14 @@ def get_audit_models(registry: orm.registry):
     return models
 
 
+def inject_model_metadata(models: t.Iterable[type[Audit]], **kwargs):
+    kwargs.setdefault("schema", "public")
+
+    for m in models:
+        for k, v in kwargs.items():
+            setattr(m, f"__audit_{k}__", v)
+
+
 def is_object_modified(obj: object):
     if isinstance(obj, Audit):
         excluded = set(obj.__class__.__audit_args__.get("exclude", []))
@@ -242,6 +279,7 @@ class PostgreSQLAudit:
         )
 
         self.pg_audit_classes = get_audit_models(self.Base.registry)
+        inject_model_metadata(self.pg_audit_classes, schema=self.schema_name)
 
         self.setup_db()
         self.attach_listeners()
@@ -249,40 +287,6 @@ class PostgreSQLAudit:
         app.extensions["postgresql-audit"] = self
 
     def fetch_activity(self, obj: type[Audit] | Audit | t.Sequence[Audit]):
-        if inspect.isclass(obj):
-            source = obj
-            relid = sa.text(f"'{source.__tablename__}'::regclass::oid")
-            jsonb = sa.func.to_jsonb(source.__table__.table_valued())
-            activity_join_cond = sa.and_(
-                self.Activity.table_name == source.__tablename__,
-                self.Activity.row_key == self.func.get_pk_values(relid, jsonb),
-            )
-
-        else:
-            objects = list(obj) if isinstance(obj, t.Sequence) else list([obj])
-            unions = []
-            for key, group in groupby(objects, key=lambda o: o.__class__):
-                relid = sa.text(f"'{key.__tablename__}'::regclass::oid")
-                jsonb = sa.func.to_jsonb(key.__table__.table_valued())
-
-                union = sa.select(
-                    sa.literal(key.__tablename__).label("table_name"),
-                    self.func.get_pk_values(relid, jsonb).label("row_key"),
-                )
-
-                for pk in key.__table__.primary_key.columns:
-                    pk_attr: sa.ColumnElement[t.Any] = getattr(key, pk.name)
-                    pks = [getattr(o, pk.name) for o in group]
-                    union = union.filter(pk_attr.in_(pks))
-
-                unions.append(union)
-
-            source = sa.union(*unions).subquery("source")
-            activity_join_cond = sa.and_(
-                self.Activity.table_name == source.c.table_name,
-                self.Activity.row_key == source.c.row_key,
-            )
-
         stmt = (
             sa.select(
                 self.Activity.id,
@@ -295,10 +299,29 @@ class PostgreSQLAudit:
                 self.Activity.old_data,
                 self.Activity.changed_data,
             )
-            .select_from(source)
-            .join(self.Activity, activity_join_cond)
+            .select_from(self.Activity)
             .join(self.Transaction)
             .order_by(self.Activity.id.desc())
         )
 
+        if inspect.isclass(obj):
+            class_ = obj
+            stmt = stmt.filter(self.Activity.table_name == class_.__tablename__)
+
+        else:
+            objects = list(obj) if isinstance(obj, t.Sequence) else list([obj])
+            conditions = []
+            for class_, objs in groupby(objects, key=lambda o: o.__class__):
+                conditions.append(
+                    sa.and_(
+                        self.Activity.table_name == class_.__tablename__,
+                        self.Activity.row_key.in_([o.__audit_rowkey__ for o in objs]),
+                    )
+                )
+
+            if len(conditions) == 1:
+                stmt = stmt.filter(conditions[0])
+            else:
+                stmt = stmt.filter(sa.or_(*conditions))
+
         return stmt

{flask_postgresql_audit-1.0.7 → flask_postgresql_audit-1.0.8}/flask_postgresql_audit/templates/get_pk_values.sql

@@ -11,7 +11,7 @@ CREATE OR REPLACE FUNCTION get_pk_values(relid oid, row_data jsonb) RETURNS json
             AND i.indisprimary;
 
         IF pk_columns IS NULL THEN
-            RETURN '{}'::jsonb;
+            RETURN '[]'::jsonb;
         END IF;
 
         SELECT jsonb_agg(row_data->>col) INTO pk_values

flask_postgresql_audit-1.0.8/flask_postgresql_audit/templates/get_table_relid.sql

@@ -0,0 +1,4 @@
+CREATE OR REPLACE FUNCTION get_table_relid(tablename text) RETURNS oid AS $$
+SELECT tablename::regclass::oid;
+$$ LANGUAGE SQL;
+