Manual rollback of committed data

Author: mesemus

pytest_invenio/database_tools.py

@@ -0,0 +1,131 @@
+import logging
+import sys
+
+from sqlalchemy import MetaData, and_, func, select
+
+logger = logging.getLogger(__name__)
+
+
+def store_database_values(engine, conn):
+    """Introspect the session, get all the tables and store their primary key values.
+
+    The result is a dict[table_name, list[pk_tuple]]
+    """
+    print("Storing database values", file=sys.stderr, flush=True)
+
+    metadata = MetaData()
+    metadata.reflect(engine)
+
+    dump = {}
+    for table_name, table in metadata.tables.items():
+        # Get primary key columns and foreign key columns
+        pk_columns = [
+            column
+            for column in table.columns
+            if column.primary_key or len(column.foreign_keys) > 0
+        ]
+
+        if not pk_columns:
+            # Skip tables without primary keys
+            continue
+
+        # Select only primary key columns
+        result = conn.execute(select(*pk_columns))
+        try:
+            dump[table_name] = [tuple(row) for row in result.fetchall()]
+        except Exception as ex:
+            raise RuntimeError(f"Could not fetch rows from table {table_name}") from ex
+
+    print("Storing database values done", file=sys.stderr, flush=True)
+
+    return dump
+
+
+def purge_database_values(engine, conn, stored_values):
+    """Delete rows that are not in the stored values."""
+
+    print("Purging database values", file=sys.stderr, flush=True)
+    metadata = MetaData()
+    metadata.reflect(engine)
+
+    # Build a list of (table_name, delete_condition) tuples
+    to_be_deleted = []
+
+    for table_name, table in metadata.tables.items():
+        stored_rows = stored_values.get(table_name, [])
+
+        # Get primary key columns and foreign key columns
+        pk_columns = [
+            column
+            for column in table.columns
+            if column.primary_key or len(column.foreign_keys) > 0
+        ]
+
+        if not pk_columns:
+            logger.warning(f"Table {table_name} has no primary key. Skipping.")
+            continue
+
+        # Convert stored rows to a set of primary key tuples for fast lookup
+        stored_pk_set = set(stored_rows)
+
+        # create a select statement that would include only rows that are not present
+        # in the stored values. It will be not (pk1 == val1 and pk2 == val2 and ...) and not (...)
+        row_matcher_conditions = []
+        for stored_pk in stored_pk_set:
+            condition = and_(
+                pk_col == pk_val for pk_col, pk_val in zip(pk_columns, stored_pk)
+            )
+            # negate the condition to match rows that are not equal
+            row_matcher_conditions.append(~condition)
+
+        if row_matcher_conditions:
+            non_matching_condition = and_(*row_matcher_conditions)
+            to_be_deleted.append(
+                (table_name, table, non_matching_condition, len(stored_pk_set))
+            )
+        else:
+            # delete everything
+            to_be_deleted.append((table_name, table, None, len(stored_pk_set)))
+
+    # Try to delete rows with retry mechanism for foreign key constraints
+    while to_be_deleted:
+        print(
+            "Purging iteration, tables to process:",
+            [table_name for table_name, _, _, _ in to_be_deleted],
+            file=sys.stderr,
+            flush=True,
+        )
+        failed_deletions = []
+
+        for table_name, table, where_condition, expected_count in to_be_deleted:
+            try:
+                delete_stmt = table.delete()
+                if where_condition is not None:
+                    delete_stmt = delete_stmt.where(where_condition)
+                conn.execute(delete_stmt)
+                conn.commit()
+                existing_count = conn.execute(
+                    select(func.count()).select_from(table)
+                ).scalar()
+                if existing_count != expected_count:
+                    logger.warning("Not all rows deleted as expected, will try again.")
+                    failed_deletions.append(
+                        (table_name, table, where_condition, expected_count)
+                    )
+            except Exception:
+                # Rollback on failure and retry in next iteration
+                conn.rollback()
+                failed_deletions.append(
+                    (table_name, table, where_condition, expected_count)
+                )
+
+        if len(failed_deletions) == len(to_be_deleted):
+            table_names = [table_name for table_name, _, _, _ in failed_deletions]
+            raise RuntimeError(
+                f"Could not delete any rows in this iteration due to foreign key cycles in tables: {table_names}"
+            )
+        else:
+            # Update the list with failed deletions for next iteration
+            to_be_deleted = failed_deletions
+
+    print("Purging database values done", file=sys.stderr, flush=True)

pytest_invenio/fixtures.py

@@ -551,6 +551,11 @@ def db(database, db_session_options):
     """
     from flask_sqlalchemy.session import Session as FlaskSQLAlchemySession
 
+    from pytest_invenio.database_tools import (
+        purge_database_values,
+        store_database_values,
+    )
+
     class PytestInvenioSession(FlaskSQLAlchemySession):
         def get_bind(self, mapper=None, clause=None, bind=None, **kwargs):
             if self.bind:
@@ -574,17 +579,31 @@ def rollback(self) -> None:
     )
     session = database._make_scoped_session(options=options)
 
-    session.begin_nested()
-
     old_session = database.session
     database.session = session
 
+    # the session.rollback() does not always clean everything, if the test
+    # used db.session.commit() and has not cleaned up after itself. We can not
+    # use nested transactions because a lot of Invenio code would need to be updated
+    # so that it is aware of the nested transaction concept. Instead, we store
+    # the database values here and purge any new rows after the test.
+    stored_values = store_database_values(database.engine, connection)
+
     yield database
 
+    print("Rollback and original connection close", file=sys.stderr, flush=True)
     session.rollback()
     connection.close()
     database.session = old_session
 
+    print("Purging database changes", file=sys.stderr, flush=True)
+    # use a brand new connection for the purge operation
+    connection = database.engine.connect()
+    connection.begin()
+    purge_database_values(database.engine, connection, stored_values)
+    connection.close()
+    print("Purged and closed", file=sys.stderr, flush=True)
+
 
 @pytest.fixture(scope="function")
 def mailbox(base_app):