feat: Use temp table for deletion

Signed-off-by: Alan Peixinho <alan.peixinho@profusion.mobi>

Author: alanpeixinho

backend/docs/prune_db command.md

@@ -16,7 +16,7 @@ Models use `DO_NOTHING` foreign keys, so the command applies manual cascade rule
 ### Optional Parameters
 
 - `--tables`: Limit pruning to specific tables (comma-separated). Valid options: `checkouts`, `builds`, `tests`. Default: all three.
-  - Cascade still applies inside the selected tables: a build matched because its checkout is old is deleted when `builds` is selected, and a test matched because its build is doomed is deleted when `tests` is selected.
+  - Cascade only drags a child when the child's parent table is also selected. For example, `--tables tests` removes only tests past the cutoff; recent tests under an old build/checkout are kept because those parents are not being pruned. With `--tables builds,tests`, an old build still drags its recent tests, but an old checkout does not drag its recent builds (checkouts are not selected).
   - Tables not listed are not deleted. For example, `--tables builds` removes old builds but leaves their tests in place. Selecting a parent without its children (e.g. only `checkouts`) can therefore leave orphaned rows.
 - `--origins`: Limit age-based pruning to specific origins (comma-separated). If omitted, any origin is considered.
   - Cascade ignores origin: once a parent row is doomed, its children are removed even if they belong to a different origin.

backend/kernelCI_app/management/commands/prune_db.py

@@ -98,70 +98,91 @@ def handle(self, *args, **options):
         dry_run = options["dry_run"]
         origins = options["origins"]
 
-        params = {"cutoff": cutoff, "batch_size": options["batch_size"]}
+        params = {"cutoff": cutoff}
         origins_condition = ""
         if origins:
-            params["origins"] = list(origins)
+            params["origins"] = origins
             origins_condition = "AND origin = ANY(%(origins)s)"
 
-        where_clauses = self._build_where_clauses(origins_condition, protect_incidents)
+        where_clauses = self._build_where_clauses(
+            origins_condition, protect_incidents, selected_tables
+        )
+        temp_tables = {t: f"prune_{t}" for t in selected_tables}
 
         with connections["default"].cursor() as cursor:
-            counts = {
-                t: self._count(cursor, t, where_clauses[t], params)
-                for t in selected_tables
-            }
-            total = sum(counts.values())
-
-            lines = [f"Rows older than {cutoff.isoformat()}:"]
-            lines += [f"* {t}:\t{counts[t]:>8}" for t in selected_tables]
-            lines += ["----------------------", f"* total:\t{total:>8}"]
-            lines.append("Note: counts include children cascaded from pruned parents.")
-            if protect_incidents:
-                lines.append("Note: rows linked to an incident are kept.")
-            self.stdout.write("\n".join(lines))
+            try:
+                # Snapshot ids while parents still exist, so child predicates can
+                # resolve them. Do all tables before deleting anything.
+                for table in selected_tables:
+                    self._materialize(
+                        cursor, table, temp_tables[table], where_clauses[table], params
+                    )
 
-            if total == 0:
-                self.stdout.write(self.style.SUCCESS("Nothing to prune."))
-                return
+                counts = {
+                    t: self._count(cursor, temp_tables[t]) for t in selected_tables
+                }
+                total = sum(counts.values())
 
-            if dry_run:
-                self.stdout.write(
-                    self.style.WARNING(
-                        "[DRY RUN] No rows deleted. Run without --dry-run to execute."
-                    )
+                lines = [f"Rows older than {cutoff.isoformat()}:"]
+                lines += [f"* {t}:\t{counts[t]:>8}" for t in selected_tables]
+                lines += ["----------------------", f"* total:\t{total:>8}"]
+                lines.append(
+                    "Note: counts include children cascaded from pruned parents."
                 )
-                return
-
-            if not options["yes"]:
-                summary = ", ".join(f"{counts[t]} {t}" for t in selected_tables)
-                try:
-                    answer = (
-                        input(f"Delete {summary} ({total} rows total)? [y/N] ")
-                        .strip()
-                        .lower()
+                if protect_incidents:
+                    lines.append("Note: rows linked to an incident are kept.")
+                self.stdout.write("\n".join(lines))
+
+                if total == 0:
+                    self.stdout.write(self.style.SUCCESS("Nothing to prune."))
+                    return
+
+                if dry_run:
+                    self.stdout.write(
+                        self.style.WARNING(
+                            "[DRY RUN] No rows deleted. Run without --dry-run to "
+                            "execute."
+                        )
                     )
-                except EOFError:
-                    answer = ""
-                if answer not in ("y", "yes"):
-                    self.stdout.write("Aborted.")
                     return
 
-            # Delete child-first (reverse of PRUNABLE_TABLES order): each batch commits
-            # on its own, so a crash mid-run leaves children already gone before their
-            # parents, never the reverse. Reordering this would risk orphans.
-            deleted = 0
-            for table in reversed(selected_tables):
-                deleted += self._batch_delete(
-                    cursor, table, where_clauses[table], params
-                )
+                if not options["yes"]:
+                    summary = ", ".join(f"{counts[t]} {t}" for t in selected_tables)
+                    try:
+                        answer = (
+                            input(f"Delete {summary} ({total} rows total)? [y/N] ")
+                            .strip()
+                            .lower()
+                        )
+                    except EOFError:
+                        answer = ""
+                    if answer not in ("y", "yes"):
+                        self.stdout.write("Aborted.")
+                        return
 
-        self.stdout.write(self.style.SUCCESS(f"Successfully pruned {deleted} rows."))
+                # Delete child-first (reverse of PRUNABLE_TABLES order): each batch
+                # commits on its own, so a crash mid-run leaves children already gone
+                # before their parents, never the reverse. Reordering this would risk
+                # orphans.
+                deleted = 0
+                for table in reversed(selected_tables):
+                    deleted += self._batch_delete(
+                        cursor, table, temp_tables[table], options["batch_size"]
+                    )
 
-    def _build_where_clauses(self, origins_condition, protect_incidents):
+                self.stdout.write(
+                    self.style.SUCCESS(f"Successfully pruned {deleted} rows.")
+                )
+            finally:
+                for temp_table in temp_tables.values():
+                    cursor.execute(f"DROP TABLE IF EXISTS {temp_table}")
+
+    def _build_where_clauses(
+        self, origins_condition, protect_incidents, selected_tables
+    ):
         """Build the per-table WHERE clauses, chaining the cascade so each child
-        matches when its parent is doomed, and appending incident protection when
-        enabled."""
+        matches when its selected parent is doomed, and appending incident protection
+        when enabled."""
         age = f"_timestamp < %(cutoff)s {origins_condition}"
 
         # A row tied to an incident is protected, along with the ancestors that would
@@ -186,32 +207,50 @@ def _build_where_clauses(self, origins_condition, protect_incidents):
         )
 
         checkout_where = age + exclusion.get("checkouts", "")
-        build_where = (
-            f"(({age}) OR checkout_id IN (SELECT id FROM checkouts WHERE {checkout_where}))"
-            + exclusion.get("builds", "")
-        )
-        test_where = (
-            f"(({age}) OR build_id IN (SELECT id FROM builds WHERE {build_where}))"
-            + exclusion.get("tests", "")
-        )
+
+        build_terms = [f"({age})"]
+        if "checkouts" in selected_tables:
+            build_terms.append(
+                f"checkout_id IN (SELECT id FROM checkouts WHERE {checkout_where})"
+            )
+        build_where = "(" + " OR ".join(build_terms) + ")" + exclusion.get("builds", "")
+
+        test_terms = [f"({age})"]
+        if "builds" in selected_tables:
+            test_terms.append(
+                f"build_id IN (SELECT id FROM builds WHERE {build_where})"
+            )
+        test_where = "(" + " OR ".join(test_terms) + ")" + exclusion.get("tests", "")
+
         return {
             "checkouts": checkout_where,
             "builds": build_where,
             "tests": test_where,
         }
 
-    def _count(self, cursor, table, where, params):
-        cursor.execute(f"SELECT COUNT(*) FROM {table} WHERE {where}", params)
+    def _materialize(self, cursor, table, temp_table, where, params):
+        """Snapshot the doomed ids into a temp table so the nested predicate runs
+        once instead of per batch."""
+        cursor.execute(f"DROP TABLE IF EXISTS {temp_table}")
+        cursor.execute(
+            f"CREATE TEMP TABLE {temp_table} AS SELECT id FROM {table} WHERE {where}",
+            params,
+        )
+
+    def _count(self, cursor, temp_table):
+        cursor.execute(f"SELECT COUNT(*) FROM {temp_table}")
         return cursor.fetchone()[0]
 
-    def _batch_delete(self, cursor, table, where, params):
-        delete_sql = (
-            f"DELETE FROM {table} WHERE id IN "
-            f"(SELECT id FROM {table} WHERE {where} LIMIT %(batch_size)s)"
+    def _batch_delete(self, cursor, table, temp_table, batch_size):
+        sql = (
+            f"WITH batch AS ("
+            f"DELETE FROM {temp_table} WHERE id IN "
+            f"(SELECT id FROM {temp_table} LIMIT %(batch_size)s) RETURNING id"
+            f") DELETE FROM {table} WHERE id IN (SELECT id FROM batch)"
         )
         deleted_total = 0
         while True:
-            cursor.execute(delete_sql, params)
+            cursor.execute(sql, {"batch_size": batch_size})
             deleted = cursor.rowcount
             if deleted == 0:
                 break

backend/kernelCI_app/tests/integrationTests/prune_db_test.py

@@ -130,29 +130,57 @@ def test_origins_cascade_ignores_child_origin():
 
 @pytest.mark.django_db
 def test_tables_tests_only():
-    """--tables limits deletion to the selected table."""
+    """--tables tests deletes only old tests; parents and a recent test under an old
+    build/checkout are kept because those parents are not being pruned."""
     checkout = CheckoutFactory(field_timestamp=_days_ago(30))
     build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(30))
-    test = TestFactory(build=build, field_timestamp=_days_ago(30))
+    old_test = TestFactory(build=build, field_timestamp=_days_ago(30))
+    recent_test = TestFactory(build=build, field_timestamp=_days_ago(1))
 
     _prune(yes=True, tables=["tests"])
 
     assert Checkouts.objects.filter(id=checkout.id).exists()
     assert Builds.objects.filter(id=build.id).exists()
-    assert not Tests.objects.filter(id=test.id).exists()
+    assert not Tests.objects.filter(id=old_test.id).exists()
+    assert Tests.objects.filter(id=recent_test.id).exists()
 
 
 @pytest.mark.django_db
 def test_tables_builds_only():
+    """--tables builds deletes only old builds; a recent build survives, the checkout is
+    untouched, and an unlisted child test is left in place (orphaned)."""
     checkout = CheckoutFactory(field_timestamp=_days_ago(30))
-    build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(30))
-    test = TestFactory(build=build, field_timestamp=_days_ago(30))
+    old_build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(30))
+    orphaned_test = TestFactory(build=old_build, field_timestamp=_days_ago(30))
+    recent_build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(1))
 
     _prune(yes=True, tables=["builds"])
 
     assert Checkouts.objects.filter(id=checkout.id).exists()
-    assert not Builds.objects.filter(id=build.id).exists()
-    assert Tests.objects.filter(id=test.id).exists()
+    assert not Builds.objects.filter(id=old_build.id).exists()
+    assert Builds.objects.filter(id=recent_build.id).exists()
+    assert Tests.objects.filter(id=orphaned_test.id).exists()
+
+
+@pytest.mark.django_db
+def test_tables_builds_tests_cascade_within_selection():
+    """With builds and tests selected, an old build still drags its recent tests, but
+    an old checkout does not drag its recent builds (checkouts not selected)."""
+    checkout = CheckoutFactory(field_timestamp=_days_ago(30))
+    old_build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(30))
+    old_build_recent_test = TestFactory(build=old_build, field_timestamp=_days_ago(1))
+    recent_build = BuildFactory(checkout=checkout, field_timestamp=_days_ago(1))
+    recent_build_recent_test = TestFactory(
+        build=recent_build, field_timestamp=_days_ago(1)
+    )
+
+    _prune(yes=True, tables=["builds", "tests"])
+
+    assert Checkouts.objects.filter(id=checkout.id).exists()
+    assert not Builds.objects.filter(id=old_build.id).exists()
+    assert not Tests.objects.filter(id=old_build_recent_test.id).exists()
+    assert Builds.objects.filter(id=recent_build.id).exists()
+    assert Tests.objects.filter(id=recent_build_recent_test.id).exists()
 
 
 @pytest.mark.django_db