diff --git a/api/tasks/remove_app_and_related_data_task.py b/api/tasks/remove_app_and_related_data_task.py index f8af48e0990..d5892df1d87 100644 --- a/api/tasks/remove_app_and_related_data_task.py +++ b/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:") diff --git a/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py index 3e9a7859030..930c5ed4cab 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_remove_app_and_related_data_task.py +++ b/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 diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 8774760f2c4..e1911e3aa10 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/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)