flask-postgresql-audit 1.0.3__tar.gz → 1.0.5__tar.gz

{flask_postgresql_audit-1.0.3 → flask_postgresql_audit-1.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flask-postgresql_audit
-Version: 1.0.3
+Version: 1.0.5
 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.3 → flask_postgresql_audit-1.0.5}/flask_postgresql_audit/__init__.py

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

{flask_postgresql_audit-1.0.3 → flask_postgresql_audit-1.0.5}/flask_postgresql_audit/alembic/entities.py

@@ -46,6 +46,15 @@ def trigger_update_factory(**context):
     )
 
 
+def get_pk_values(**context):
+    sql = parse_template("get_pk_values.sql", **context)
+    return PGFunction(
+        schema=context["schema_name"],
+        signature="get_pk_values(relid oid, row_data jsonb)",
+        definition=sql["definition"],
+    )
+
+
 def get_setting_factory(**context):
     sql = parse_template("get_setting.sql", **context)
     return PGFunction(

{flask_postgresql_audit-1.0.3 → flask_postgresql_audit-1.0.5}/flask_postgresql_audit/alembic/setup.py

@@ -62,6 +62,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_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.3 → flask_postgresql_audit-1.0.5}/flask_postgresql_audit/base.py

@@ -22,6 +22,7 @@ class ActivityBase:
     native_transaction_id: OMap[int] = mapped_column(BigInteger, index=True)
     verb: OMap[str] = mapped_column(Text)
 
+    row_key: OMap[list[str]] = mapped_column(JSONB, default=[], server_default="[]")
     old_data: OMap[dict] = mapped_column(JSONB, default={}, server_default="{}")
     changed_data: OMap[dict] = mapped_column(JSONB, default={}, server_default="{}")
 

flask_postgresql_audit-1.0.5/flask_postgresql_audit/extensions/alembic_utils/__init__.py

@@ -0,0 +1,9 @@
+import alembic_utils
+from alembic_utils.experimental import collect_subclasses
+from alembic_utils.replaceable_entity import ReplaceableEntity
+
+from .pg_aggregate import PGAggregate
+
+original_types = set(collect_subclasses(alembic_utils, ReplaceableEntity))
+extension_types = {PGAggregate}
+allowed_types = original_types | extension_types

flask_postgresql_audit-1.0.5/flask_postgresql_audit/extensions/alembic_utils/pg_aggregate.py

@@ -0,0 +1,119 @@
+import re
+import typing as t
+
+import sqlalchemy as sa
+from alembic_utils.exceptions import SQLParseFailure
+from alembic_utils.replaceable_entity import ReplaceableEntity
+from alembic_utils.statement import strip_terminating_semicolon
+from sqlalchemy.orm import Session
+
+T = t.TypeVar("T", bound="ReplaceableEntity")
+
+
+class PGAggregate(ReplaceableEntity):
+    def __init__(self, schema: str, signature: str, definition: str):
+        super().__init__(schema, signature, definition)
+        self.definition: str = strip_terminating_semicolon(definition)
+
+    @property
+    def type_(self):
+        return "aggregate"
+
+    @property
+    def literal_signature(self):
+        name, remainder = self.signature.split("(", 1)
+        return '"' + name.strip() + '"(' + remainder
+
+    @classmethod
+    def from_sql(cls, sql: str) -> "PGAggregate":
+        pattern = r"^CREATE\s+AGGREGATE\s+(?:(?P<schema>\w+)\.)?(?P<signature>.+)\s+(?P<definition>\(.*\))$"
+        result = re.search(pattern, sql, re.IGNORECASE)
+        if result is not None and isinstance(result, re.Match):
+            raw_signature: str = result["signature"]
+            signature = (
+                "".join(raw_signature.split('"', 2))
+                if raw_signature.startswith('"')
+                else raw_signature
+            )
+            return cls(
+                schema=result["schema"],
+                signature=signature,
+                definition=result["definition"],
+            )
+        raise SQLParseFailure(f'Failed to parse SQL into PGAggregate """{sql}"""')
+
+    def to_sql_statement_create(self, *, replace: bool = False):
+        return sa.text(
+            f"{'CREATE OR REPLACE' if replace else 'CREATE'} AGGREGATE "
+            f"{self.literal_schema}.{self.literal_signature} {self.definition}"
+        )
+
+    def to_sql_statement_create_or_replace(self):
+        return self.to_sql_statement_create(replace=True)
+
+    def to_sql_statement_drop(self, cascade=False):
+        return sa.text(
+            f"DROP AGGREGATE {self.literal_schema}.{self.literal_signature}"
+            + (" CASCADE" if cascade else "")
+        )
+
+    @classmethod
+    def from_database(cls, sess: Session, schema: str) -> "list[PGAggregate]":
+        PG_GTE_11 = """
+            and p.prokind = 'a'
+        """
+
+        # NOTE: not tested
+        # PG_LT_11 = """
+        #     and p.proisagg
+        #     and not p.proiswindow
+        # """
+
+        # Retrieve the postgres server version e.g. 90603 for 9.6.3 or 120003 for 12.3
+        pg_version_str = sess.execute(sa.text("show server_version_num")).scalar_one()
+        pg_version = int(pg_version_str)
+
+        sql = sa.text(
+            """
+            WITH extension_functions AS (
+                SELECT objid AS extension_function_oid, classid
+                FROM pg_depend
+                WHERE deptype='e' -- depends on an extension
+                    AND classid = 'pg_proc'::regclass -- is a proc/function
+            )
+
+            SELECT  n.nspname AS function_schema
+                ,   p.proname AS function_name
+                ,   pg_get_function_arguments(p.oid) AS function_arguments
+                ,   format('CREATE AGGREGATE %s (SFUNC = %s, STYPE = %s%s%s%s%s)'
+                    , format('%s.%s(%s)', n.nspname, p.proname, pg_get_function_arguments(p.oid))
+                    , r.aggtransfn
+                    , r.aggtranstype::regtype
+                    , ', SORTOP = '    || NULLIF(r.aggsortop, 0)::regoper
+                    , ', INITCOND = '  || quote_nullable(r.agginitval)
+                    , ', FINALFUNC = ' || NULLIF(r.aggfinalfn, 0)
+                    , CASE WHEN r.aggfinalextra THEN ', FINALFUNC_EXTRA' END
+            --         add more to cover special cases like moving-aggregate etc.
+                    ) AS create_statement
+                ,   t.typname AS return_type
+                ,   l.lanname AS function_language
+            FROM pg_aggregate r
+                LEFT JOIN pg_proc p ON p.oid = aggfnoid
+                LEFT JOIN pg_namespace n ON p.pronamespace = n.oid
+                LEFT JOIN pg_language l ON p.prolang = l.oid
+                LEFT JOIN pg_type t ON t.oid = p.prorettype
+                LEFT JOIN extension_functions ef ON p.oid = ef.extension_function_oid
+            WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') 
+                AND ef.extension_function_oid IS NULL -- Filter out functions from extensions
+                AND n.nspname = :schema
+            """
+            + (PG_GTE_11 if pg_version >= 110000 else "")
+        ).bindparams(schema=schema)
+
+        rows = sess.execute(sql)
+        db_aggregates = [cls.from_sql(x[3]) for x in rows]
+
+        for agg in db_aggregates:
+            assert agg is not None
+
+        return db_aggregates

{flask_postgresql_audit-1.0.3 → flask_postgresql_audit-1.0.5}/flask_postgresql_audit/templates/create_activity.sql

@@ -20,7 +20,7 @@ BEGIN
     IF (TG_OP = 'UPDATE') THEN
         INSERT INTO ${schema_prefix}pga_activity(
             id, schema_name, table_name, relid, issued_at, native_transaction_id,
-            verb, old_data, changed_data, transaction_id)
+            verb, row_key, old_data, changed_data, transaction_id)
         SELECT
             nextval('${schema_prefix}pga_activity_id_seq') as id,
             TG_TABLE_SCHEMA::text AS schema_name,
@@ -29,6 +29,7 @@ BEGIN
             statement_timestamp() AT TIME ZONE 'UTC' AS issued_at,
             txid_current() AS native_transaction_id,
             LOWER(TG_OP) AS verb,
+            ${schema_prefix}get_pk_values(TG_RELID, new_data) AS row_key,
             old_data - excluded_cols AS old_data,
             ${schema_prefix}jsonb_subtract(new_data, old_data) - excluded_cols AS changed_data,
             _transaction_id AS transaction_id
@@ -48,7 +49,7 @@ BEGIN
     ELSIF (TG_OP = 'INSERT') THEN
         INSERT INTO ${schema_prefix}pga_activity(
             id, schema_name, table_name, relid, issued_at, native_transaction_id,
-            verb, old_data, changed_data, transaction_id)
+            verb, row_key, old_data, changed_data, transaction_id)
         SELECT
             nextval('${schema_prefix}pga_activity_id_seq') as id,
             TG_TABLE_SCHEMA::text AS schema_name,
@@ -57,6 +58,7 @@ BEGIN
             statement_timestamp() AT TIME ZONE 'UTC' AS issued_at,
             txid_current() AS native_transaction_id,
             LOWER(TG_OP) AS verb,
+            ${schema_prefix}get_pk_values(TG_RELID, to_jsonb(new_table)) AS row_key,
             '{}'::jsonb AS old_data,
             row_to_json(new_table.*)::jsonb - excluded_cols AS changed_data,
             _transaction_id AS transaction_id
@@ -64,7 +66,7 @@ BEGIN
     ELSEIF TG_OP = 'DELETE' THEN
         INSERT INTO ${schema_prefix}pga_activity(
             id, schema_name, table_name, relid, issued_at, native_transaction_id,
-            verb, old_data, changed_data, transaction_id)
+            verb, row_key, old_data, changed_data, transaction_id)
         SELECT
             nextval('${schema_prefix}pga_activity_id_seq') as id,
             TG_TABLE_SCHEMA::text AS schema_name,
@@ -73,6 +75,7 @@ BEGIN
             statement_timestamp() AT TIME ZONE 'UTC' AS issued_at,
             txid_current() AS native_transaction_id,
             LOWER(TG_OP) AS verb,
+            ${schema_prefix}get_pk_values(TG_RELID, to_jsonb(old_table)) AS row_key,
             row_to_json(old_table.*)::jsonb - excluded_cols AS old_data,
             '{}'::jsonb AS changed_data,
             _transaction_id AS transaction_id

flask_postgresql_audit-1.0.5/flask_postgresql_audit/templates/get_pk_values.sql

@@ -0,0 +1,22 @@
+CREATE OR REPLACE FUNCTION get_pk_values(relid oid, row_data jsonb) RETURNS jsonb AS $$
+    DECLARE
+        pk_columns text [];
+        pk_values jsonb;
+    BEGIN
+        SELECT array_agg(a.attname ORDER BY a.attnum) INTO pk_columns
+        FROM pg_index i
+            JOIN pg_attribute a ON a.attrelid = i.indrelid
+                AND a.attnum = ANY(i.indkey)
+        WHERE i.indrelid = relid
+            AND i.indisprimary;
+
+        IF pk_columns IS NULL THEN
+            RETURN '{}'::jsonb;
+        END IF;
+
+        SELECT jsonb_agg(row_data->>col) INTO pk_values
+        FROM unnest(pk_columns) AS col;
+
+        RETURN pk_values;
+    END;
+$$ LANGUAGE plpgsql;

{flask_postgresql_audit-1.0.3 → flask_postgresql_audit-1.0.5}/tests/test_defaults.py

@@ -9,6 +9,7 @@ class TestCustomSchemaactivityCreation:
     def test_insert(self, user):
         stmt = sa.select(audit.Activity).limit(1)
         if activity := db.session.scalar(stmt):
+            assert activity.row_key == ["1"]
             assert activity.old_data == {}
             assert activity.changed_data == {"id": user.id, "name": "John", "age": 15}
             assert activity.table_name == "user"