refactor(BA-5936): apply Copilot review feedback

- Wrap all DELETEs (cascade, RBAC, parent) in a single try/except for
  IntegrityError so cascade failures also surface as
  RepositoryIntegrityError instead of raw SQLAlchemy errors.
- Drop PrunerSpec.returning_id() — derive the parent PK column directly
  from row_class().__table__.primary_key, mirroring the Purger pattern.
  Reject composite-PK tables with UnsupportedCompositePrimaryKeyError.
- Update CascadeChild docstring to reflect the materialized-id-list
  approach (the previous SQL-subquery wording was stale).
- Tighten test_no_cascade_with_fk_violation_raises to expect
  ForeignKeyViolationError, also covering the parse_integrity_error path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: fregataa

src/ai/backend/manager/repositories/base/pruner.py

@@ -9,6 +9,7 @@
 import sqlalchemy as sa
 
 from ai.backend.common.data.permission.types import EntityType
+from ai.backend.manager.errors.repository import UnsupportedCompositePrimaryKeyError
 from ai.backend.manager.models.base import Base
 from ai.backend.manager.models.rbac_models.association_scopes_entities import (
     AssociationScopesEntitiesRow,
@@ -26,11 +27,14 @@
 class CascadeChild(ABC):
     """A child table whose rows must be deleted before the parent's prune.
 
-    Used for simple FK cascades. Each cascade DELETE runs as::
+    Used for simple FK cascades. ``execute_pruner`` first locks and
+    materializes the parent target IDs once, then issues each cascade
+    DELETE as::
 
-        DELETE FROM <row_class> WHERE <parent_id_column>
-            IN (SELECT <parent pk> FROM <parent>
-                WHERE <prune_condition AND conditions>)
+        DELETE FROM <row_class> WHERE <parent_id_column> IN (<target_ids>)
+
+    where ``<target_ids>`` is the list returned from the single
+    ``SELECT pk FOR UPDATE`` against the parent table.
 
     Polymorphic / cross-cutting cleanups (e.g., RBAC associations) are not
     handled here — see :meth:`PrunerSpec.entity_type` for that.
@@ -94,18 +98,13 @@ class PrunerSpec[TRow: Base](ABC):
     def row_class(cls) -> type[TRow]:
         """ORM Row class for the parent entity table.
 
-        Example:
-            return SessionRow
-        """
-        raise NotImplementedError
-
-    @classmethod
-    @abstractmethod
-    def returning_id(cls) -> Any:
-        """Primary-key column for the parent's ``DELETE ... RETURNING``.
+        The single-column primary key is derived from
+        ``row_class().__table__.primary_key`` by ``execute_pruner``;
+        composite-PK tables are rejected with
+        :class:`UnsupportedCompositePrimaryKeyError`.
 
         Example:
-            return SessionRow.id
+            return SessionRow
         """
         raise NotImplementedError
 
@@ -175,11 +174,19 @@ async def execute_pruner[TRow: Base](
         PrunerResult with the count and PK list of deleted parent rows.
 
     Raises:
-        RepositoryIntegrityError: If any DELETE violates a database constraint.
+        UnsupportedCompositePrimaryKeyError: If the parent table has a
+            composite primary key.
+        RepositoryIntegrityError: If any DELETE (cascade, RBAC, or parent)
+            violates a database constraint.
     """
     cls = type(spec)
     table = cls.row_class().__table__
-    pk_col = cls.returning_id()
+    pk_columns = list(table.primary_key.columns)
+    if len(pk_columns) != 1:
+        raise UnsupportedCompositePrimaryKeyError(
+            f"PrunerSpec only supports single-column primary keys (table: {table.name})",
+        )
+    pk_col = pk_columns[0]
 
     where = cls.prune_condition()
     for f in spec.conditions:
@@ -190,24 +197,24 @@ async def execute_pruner[TRow: Base](
     if not target_ids:
         return PrunerResult(count=0, ids=[])
 
-    for child in spec.cascade:
-        ccls = type(child)
-        cascade_table = ccls.row_class().__table__
-        await db_sess.execute(
-            sa.delete(cascade_table).where(ccls.parent_id_column().in_(target_ids))
-        )
-
     rbac_entity_type = cls.entity_type()
-    if spec.cascade_rbac and rbac_entity_type is not None:
-        await db_sess.execute(
-            sa.delete(AssociationScopesEntitiesRow).where(
-                AssociationScopesEntitiesRow.entity_type == rbac_entity_type,
-                AssociationScopesEntitiesRow.entity_id.in_([str(i) for i in target_ids]),
+    try:
+        for child in spec.cascade:
+            ccls = type(child)
+            cascade_table = ccls.row_class().__table__
+            await db_sess.execute(
+                sa.delete(cascade_table).where(ccls.parent_id_column().in_(target_ids))
             )
-        )
 
-    stmt = sa.delete(table).where(pk_col.in_(target_ids)).returning(pk_col)
-    try:
+        if spec.cascade_rbac and rbac_entity_type is not None:
+            await db_sess.execute(
+                sa.delete(AssociationScopesEntitiesRow).where(
+                    AssociationScopesEntitiesRow.entity_type == rbac_entity_type,
+                    AssociationScopesEntitiesRow.entity_id.in_([str(i) for i in target_ids]),
+                )
+            )
+
+        stmt = sa.delete(table).where(pk_col.in_(target_ids)).returning(pk_col)
         deleted = list((await db_sess.scalars(stmt)).all())
     except sa.exc.IntegrityError as e:
         raise parse_integrity_error(e) from e

tests/unit/manager/repositories/base/test_pruner.py

@@ -14,6 +14,7 @@
 from sqlalchemy.dialects.postgresql import UUID as PGUUID
 
 from ai.backend.common.data.permission.types import EntityType, ScopeType
+from ai.backend.manager.errors.repository import ForeignKeyViolationError
 from ai.backend.manager.models.base import Base
 from ai.backend.manager.models.rbac_models.association_scopes_entities import (
     AssociationScopesEntitiesRow,
@@ -77,10 +78,6 @@ class TerminatedTestParentPrunerSpec(PrunerSpec[PrunerTestParentRow]):
     def row_class(cls) -> type[PrunerTestParentRow]:
         return PrunerTestParentRow
 
-    @classmethod
-    def returning_id(cls) -> Any:
-        return PrunerTestParentRow.id
-
     @classmethod
     def prune_condition(cls) -> sa.ColumnElement[bool]:
         return PrunerTestParentRow.status == "terminated"
@@ -325,10 +322,15 @@ async def test_cascade_skipped_for_non_terminal_parents(
     async def test_no_cascade_with_fk_violation_raises(
         self, parent_child_tables: ExtendedAsyncSAEngine
     ) -> None:
-        """Without the cascade, FK constraint blocks the parent DELETE."""
+        """Without the cascade, FK constraint blocks the parent DELETE.
+
+        Also verifies that ``execute_pruner`` translates the SQLAlchemy
+        ``IntegrityError`` into ``ForeignKeyViolationError`` via
+        ``parse_integrity_error``.
+        """
         await self._seed_with_children(parent_child_tables)
 
-        with pytest.raises(Exception):
+        with pytest.raises(ForeignKeyViolationError):
             async with parent_child_tables.begin_session() as db_sess:
                 spec = TerminatedTestParentPrunerSpec()  # no cascade
                 await execute_pruner(db_sess, spec)