fix(api): delete offload objects outside the DB transaction (#42258)

api/tasks/remove_app_and_related_data_task.py

@@ -513,7 +513,9 @@ def delete_draft_variables_batch(app_id: str, batch_size: int = 1000) -> int:
     total_files_deleted = 0
 
     while True:
-        with session_factory.create_session() as session, session.begin():
+        # Read the batch in a transaction of its own: the offload cleanup in between
+        # must run with no transaction open.
+        with session_factory.create_session() as session:
             # Get a batch of draft variable IDs along with their file_ids
             query_sql = """
                 SELECT id, file_id FROM workflow_draft_variables
@@ -523,18 +525,21 @@ def delete_draft_variables_batch(app_id: str, batch_size: int = 1000) -> int:
             result = session.execute(sa.text(query_sql), {"app_id": app_id, "batch_size": batch_size})
 
             rows = list(result)
-            if not rows:
-                break
 
-            draft_var_ids = [row[0] for row in rows]
-            file_ids = [row[1] for row in rows if row[1] is not None]
+        if not rows:
+            break
 
-            # Clean up associated Offload data first
-            if file_ids:
-                files_deleted = _delete_draft_variable_offload_data(session, file_ids)
-                total_files_deleted += files_deleted
+        draft_var_ids = [row[0] for row in rows]
+        file_ids = [row[1] for row in rows if row[1] is not None]
 
-            # Delete the draft variables
+        # Clean up associated Offload data first, outside any transaction, so the rows
+        # remain visible if the storage cleanup has to be retried.
+        if file_ids:
+            files_deleted = _delete_draft_variable_offload_data(file_ids)
+            total_files_deleted += files_deleted
+
+        # Delete the draft variables
+        with session_factory.create_session() as session, session.begin():
             delete_sql = """
                 DELETE FROM workflow_draft_variables
                 WHERE id IN :ids
@@ -558,7 +563,7 @@ def delete_draft_variables_batch(app_id: str, batch_size: int = 1000) -> int:
     return total_deleted
 
 
-def _delete_draft_variable_offload_data(session, file_ids: list[str]) -> int:
+def _delete_draft_variable_offload_data(file_ids: list[str]) -> int:
     """
     Delete Offload data associated with WorkflowDraftVariable file_ids.
 
@@ -568,8 +573,10 @@ def _delete_draft_variable_offload_data(session, file_ids: list[str]) -> int:
     3. Deletes UploadFile records
     4. Deletes WorkflowDraftVariableFile records
 
+    Object storage is external: the storage deletes below run outside any database
+    transaction (see api/AGENTS.md).
+
     Args:
-        session: Database connection
         file_ids: List of WorkflowDraftVariableFile IDs
 
     Returns:
@@ -590,8 +597,9 @@ def _delete_draft_variable_offload_data(session, file_ids: list[str]) -> int:
                              JOIN upload_files uf ON wdvf.upload_file_id = uf.id
                     WHERE wdvf.id IN :file_ids \
                     """
-        result = session.execute(sa.text(query_sql), {"file_ids": tuple(file_ids)})
-        file_records = list(result)
+        with session_factory.create_session() as session:
+            result = session.execute(sa.text(query_sql), {"file_ids": tuple(file_ids)})
+            file_records = [(row[0], row[1], row[2]) for row in result]
 
         # Delete from object storage and collect upload file IDs
         upload_file_ids = []
@@ -605,22 +613,23 @@ def _delete_draft_variable_offload_data(session, file_ids: list[str]) -> int:
                 # Continue with database cleanup even if storage deletion fails
                 upload_file_ids.append(upload_file_id)
 
-        # Delete UploadFile records
-        if upload_file_ids:
-            delete_upload_files_sql = """
-                                      DELETE \
-                                      FROM upload_files
-                                      WHERE id IN :upload_file_ids \
-                                      """
-            session.execute(sa.text(delete_upload_files_sql), {"upload_file_ids": tuple(upload_file_ids)})
+        with session_factory.create_session() as session, session.begin():
+            # Delete UploadFile records
+            if upload_file_ids:
+                delete_upload_files_sql = """
+                                          DELETE \
+                                          FROM upload_files
+                                          WHERE id IN :upload_file_ids \
+                                          """
+                session.execute(sa.text(delete_upload_files_sql), {"upload_file_ids": tuple(upload_file_ids)})
 
-        # Delete WorkflowDraftVariableFile records
-        delete_variable_files_sql = """
-                                    DELETE \
-                                    FROM workflow_draft_variable_files
-                                    WHERE id IN :file_ids \
-                                    """
-        session.execute(sa.text(delete_variable_files_sql), {"file_ids": tuple(file_ids)})
+            # Delete WorkflowDraftVariableFile records
+            delete_variable_files_sql = """
+                                        DELETE \
+                                        FROM workflow_draft_variable_files
+                                        WHERE id IN :file_ids \
+                                        """
+            session.execute(sa.text(delete_variable_files_sql), {"file_ids": tuple(file_ids)})
 
     except Exception:
         logging.exception("Error deleting draft variable offload data:")

api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py

@@ -6,7 +6,6 @@ import pytest
 from sqlalchemy import delete, func, select
 from sqlalchemy.orm import Session
 
-from core.db.session_factory import session_factory
 from extensions.storage.storage_type import StorageType
 from graphon.variables.segments import StringSegment
 from graphon.variables.types import SegmentType
@@ -168,7 +167,7 @@ class TestDeleteDraftVariablesBatch:
 
         assert result == 30
         mock_offload_cleanup.assert_called_once()
-        _, called_file_ids = mock_offload_cleanup.call_args.args
+        (called_file_ids,) = mock_offload_cleanup.call_args.args
         assert {str(file_id) for file_id in called_file_ids} == {str(file_id) for file_id in file_id_by_index.values()}
         info_records = [record for record in caplog.records if record.levelno == logging.INFO]
         assert len(info_records) == 2
@@ -186,8 +185,7 @@ class TestDeleteDraftVariableOffloadData:
         upload_file_keys = [upload_file.key for upload_file in offload_data["upload_files"]]
         upload_file_ids = [upload_file.id for upload_file in offload_data["upload_files"]]
 
-        with session_factory.create_session() as session, session.begin():
-            result = _delete_draft_variable_offload_data(session, file_ids)
+        result = _delete_draft_variable_offload_data(file_ids)
 
         assert result == 3
         expected_storage_calls = [call(storage_key) for storage_key in upload_file_keys]
@@ -218,8 +216,7 @@ class TestDeleteDraftVariableOffloadData:
         mock_storage.delete.side_effect = [Exception("Storage error"), None]
 
         with caplog.at_level(logging.ERROR):
-            with session_factory.create_session() as session, session.begin():
-                result = _delete_draft_variable_offload_data(session, file_ids)
+            result = _delete_draft_variable_offload_data(file_ids)
 
         assert result == 1
         assert f"Failed to delete storage object {storage_keys[0]}" in caplog.text

api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py

@@ -1,4 +1,6 @@
 import logging
+from collections.abc import Iterator, Sequence
+from contextlib import AbstractContextManager
 from datetime import UTC, datetime
 from unittest.mock import MagicMock, call, patch
 from uuid import uuid4
@@ -143,30 +145,175 @@ class TestDeleteDraftVariablesBatch:
         mock_batch_delete.assert_called_once_with(app_id, batch_size=1000)
 
 
+class _FakeResult:
+    """Minimal stand-in for a SQLAlchemy ``Result``."""
+
+    def __init__(self, rows: Sequence[Sequence[object]] = (), rowcount: int = 0) -> None:
+        self._rows = [tuple(row) for row in rows]
+        self.rowcount = rowcount
+
+    def __iter__(self) -> Iterator[Sequence[object]]:
+        return iter(self._rows)
+
+
+class _FakeSessionFactory:
+    """Session factory stub that counts open sessions and hands out queued results.
+
+    Lets the tests assert that object storage cleanup runs with no session open.
+    """
+
+    def __init__(self, results: list[_FakeResult] | None = None) -> None:
+        self._results = list(results or [])
+        self.open_count = 0
+        self.open_count_at_storage_delete: list[int] = []
+
+    def create_session(self) -> AbstractContextManager[MagicMock]:
+        factory = self
+
+        class _TransactionContext:
+            def __init__(self, session: MagicMock) -> None:
+                self._session = session
+
+            def __enter__(self) -> MagicMock:
+                factory.open_count += 1
+                return self._session
+
+            def __exit__(self, *_exc_info: object) -> bool:
+                factory.open_count -= 1
+                return False
+
+        class _SessionContext:
+            def __enter__(self) -> MagicMock:
+                factory.open_count += 1
+                session = MagicMock()
+                session.execute.side_effect = factory.next_result
+                session.begin.return_value = _TransactionContext(session)
+                return session
+
+            def __exit__(self, *_exc_info: object) -> bool:
+                factory.open_count -= 1
+                return False
+
+        return _SessionContext()
+
+    def next_result(self, *_args: object, **_kwargs: object) -> _FakeResult:
+        if self._results:
+            return self._results.pop(0)
+        return _FakeResult([])
+
+
 class TestDeleteDraftVariableOffloadData:
     """Test the Offload data cleanup functionality."""
 
     def test_delete_draft_variable_offload_data_empty_file_ids(self):
-        """Test handling of empty file_ids list."""
-        mock_conn = MagicMock()
-
-        result = _delete_draft_variable_offload_data(mock_conn, [])
+        """An empty file_ids list must not open a session or touch object storage."""
+        with (
+            patch("tasks.remove_app_and_related_data_task.session_factory") as mock_session_factory,
+            patch("extensions.ext_storage.storage") as mock_storage,
+        ):
+            result = _delete_draft_variable_offload_data([])
 
         assert result == 0
-        mock_conn.execute.assert_not_called()
+        mock_session_factory.create_session.assert_not_called()
+        mock_storage.delete.assert_not_called()
 
     def test_delete_draft_variable_offload_data_database_failure(self, caplog: pytest.LogCaptureFixture):
         """Test handling of database operation failures."""
-        mock_conn = MagicMock()
-        file_ids = ["file-1"]
-        mock_conn.execute.side_effect = Exception("Database error")
+        mock_session_factory = MagicMock()
+        mock_session_factory.create_session.return_value.__enter__.return_value.execute.side_effect = Exception(
+            "Database error"
+        )
 
-        with caplog.at_level(logging.ERROR):
-            result = _delete_draft_variable_offload_data(mock_conn, file_ids)
+        with patch("tasks.remove_app_and_related_data_task.session_factory", mock_session_factory):
+            with caplog.at_level(logging.ERROR):
+                result = _delete_draft_variable_offload_data(["file-1"])
 
         assert result == 0
         assert "Error deleting draft variable offload data:" in caplog.text
 
+    def test_storage_delete_runs_with_no_open_transaction(self, monkeypatch: pytest.MonkeyPatch):
+        """Object storage cleanup must not run while a database transaction is open."""
+        factory = _FakeSessionFactory([_FakeResult([("vf-1", "key-1", "uf-1"), ("vf-2", "key-2", "uf-2")])])
+        mock_storage = MagicMock()
+        mock_storage.delete.side_effect = lambda _key: factory.open_count_at_storage_delete.append(factory.open_count)
+
+        monkeypatch.setattr(remove_app_task_module, "session_factory", factory)
+        monkeypatch.setattr("extensions.ext_storage.storage", mock_storage)
+
+        result = _delete_draft_variable_offload_data(["vf-1", "vf-2"])
+
+        assert result == 2
+        assert mock_storage.delete.call_count == 2
+        assert factory.open_count_at_storage_delete == [0, 0]
+        assert factory.open_count == 0
+
+    def test_storage_failure_is_logged_and_not_counted(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ):
+        """A failing object is logged and not counted, and the DB cleanup still runs."""
+        factory = _FakeSessionFactory([_FakeResult([("vf-1", "key-1", "uf-1"), ("vf-2", "key-2", "uf-2")])])
+        mock_storage = MagicMock()
+        mock_storage.delete.side_effect = [Exception("Storage error"), None]
+
+        monkeypatch.setattr(remove_app_task_module, "session_factory", factory)
+        monkeypatch.setattr("extensions.ext_storage.storage", mock_storage)
+
+        with caplog.at_level(logging.ERROR):
+            result = _delete_draft_variable_offload_data(["vf-1", "vf-2"])
+
+        assert result == 1
+        assert "Failed to delete storage object key-1" in caplog.text
+        assert factory.open_count == 0
+
+    def test_file_ids_with_no_matching_rows_delete_nothing(self, monkeypatch: pytest.MonkeyPatch):
+        """file_ids that no longer resolve to rows must not raise or touch storage."""
+        factory = _FakeSessionFactory([_FakeResult([])])
+        mock_storage = MagicMock()
+
+        monkeypatch.setattr(remove_app_task_module, "session_factory", factory)
+        monkeypatch.setattr("extensions.ext_storage.storage", mock_storage)
+
+        result = _delete_draft_variable_offload_data(["vf-gone"])
+
+        assert result == 0
+        mock_storage.delete.assert_not_called()
+        assert factory.open_count == 0
+
+
+class TestDeleteDraftVariablesBatchTransactionScope:
+    """The batch loop must not hold a transaction while object storage is cleaned up."""
+
+    def test_offload_cleanup_is_called_with_no_session_open(self, monkeypatch: pytest.MonkeyPatch):
+        factory = _FakeSessionFactory([_FakeResult([("var-1", "vf-1")]), _FakeResult([], rowcount=1)])
+        observed_open_counts: list[int] = []
+
+        def _record_cleanup(file_ids: list[str]) -> int:
+            observed_open_counts.append(factory.open_count)
+            return len(file_ids)
+
+        monkeypatch.setattr(remove_app_task_module, "session_factory", factory)
+        monkeypatch.setattr(remove_app_task_module, "_delete_draft_variable_offload_data", _record_cleanup)
+
+        total_deleted = delete_draft_variables_batch("app-1", batch_size=10)
+
+        assert total_deleted == 1
+        assert observed_open_counts == [0]
+        assert factory.open_count == 0
+
+    def test_rows_without_offloaded_files_skip_cleanup(self, monkeypatch: pytest.MonkeyPatch):
+        """A batch whose rows carry no file_id must not invoke the offload cleanup."""
+        factory = _FakeSessionFactory([_FakeResult([("var-1", None)]), _FakeResult([], rowcount=1)])
+        cleanup = MagicMock(return_value=0)
+
+        monkeypatch.setattr(remove_app_task_module, "session_factory", factory)
+        monkeypatch.setattr(remove_app_task_module, "_delete_draft_variable_offload_data", cleanup)
+
+        total_deleted = delete_draft_variables_batch("app-1", batch_size=10)
+
+        assert total_deleted == 1
+        cleanup.assert_not_called()
+        assert factory.open_count == 0
+
 
 class TestDeleteWorkflowArchiveLogs:
     @pytest.mark.parametrize("sqlite_session", [(WorkflowArchiveLog,)], indirect=True)