Wrap schema wiping in a transaction to avoid orphaned tables in the immune schema (#461)

Author: droserasprout

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning].
 
 - codegen: Fixed invalid `models.py` template.
 - context: Do not wrap known exceptions with `CallbackError`
+- database: Wrap schema wiping in a transaction to avoid orphaned tables in the immune schema.
 - hasura: Fixed processing M2M relations.
 - sentry: Fixed "invalid value `environment`" error.
 - sentry: Ignore events from project callbacks when `crash_reporting` is enabled.

src/dipdup/cli.py

@@ -479,7 +479,7 @@ async def schema_wipe(ctx, immune: bool, force: bool) -> None:
         if isinstance(config.database, PostgresDatabaseConfig):
             await wipe_schema(
                 conn=conn,
-                name=config.database.schema_name,
+                schema_name=config.database.schema_name,
                 # NOTE: Don't be confused by the name of `--immune` flag, we want to drop all tables if it's set.
                 immune_tables=config.database.immune_tables if not immune else set(),
             )

src/dipdup/context.py

@@ -207,7 +207,11 @@ async def reindex(self, reason: Optional[Union[str, ReindexingReason]] = None, *
         elif action == ReindexingAction.wipe:
             conn = get_connection()
             if isinstance(self.config.database, PostgresDatabaseConfig):
-                await wipe_schema(conn, self.config.database.schema_name, self.config.database.immune_tables)
+                await wipe_schema(
+                    conn=conn,
+                    schema_name=self.config.database.schema_name,
+                    immune_tables=self.config.database.immune_tables,
+                )
             else:
                 await Tortoise._drop_databases()
             await self.restart()

src/dipdup/utils/database.py

@@ -164,22 +164,29 @@ async def truncate_schema(conn: BaseDBAsyncClient, name: str) -> None:
     await conn.execute_script(f"SELECT truncate_schema('{name}')")
 
 
-async def wipe_schema(conn: BaseDBAsyncClient, name: str, immune_tables: Set[str]) -> None:
+async def wipe_schema(
+    conn: BaseDBAsyncClient,
+    schema_name: str,
+    immune_tables: Set[str],
+) -> None:
+    """Truncate schema preserving immune tables. Executes in a transaction"""
     if isinstance(conn, SqliteClient):
         raise NotImplementedError
 
-    immune_schema_name = f'{name}_immune'
-    if immune_tables:
-        await create_schema(conn, immune_schema_name)
-        for table in immune_tables:
-            await move_table(conn, table, name, immune_schema_name)
+    immune_schema_name = f'{schema_name}_immune'
 
-    await truncate_schema(conn, name)
+    async with conn._in_transaction() as conn:
+        if immune_tables:
+            await create_schema(conn, immune_schema_name)
+            for table in immune_tables:
+                await move_table(conn, table, schema_name, immune_schema_name)
 
-    if immune_tables:
-        for table in immune_tables:
-            await move_table(conn, table, immune_schema_name, name)
-        await drop_schema(conn, immune_schema_name)
+        await truncate_schema(conn, schema_name)
+
+        if immune_tables:
+            for table in immune_tables:
+                await move_table(conn, table, immune_schema_name, schema_name)
+            await drop_schema(conn, immune_schema_name)
 
 
 async def drop_schema(conn: BaseDBAsyncClient, name: str) -> None:
@@ -199,8 +206,9 @@ async def move_table(conn: BaseDBAsyncClient, name: str, schema: str, new_schema
 
 def prepare_models(package: Optional[str]) -> None:
     """Prepare TortoiseORM models to use with DipDup.
-    Generate missing table names, validate models, increase decimal precision.
+    Generate missing table names, validate models, increase decimal precision if needed.
     """
+    # NOTE: Circular imports
     from dipdup.models import Model
 
     decimal_context = decimal.getcontext()