From 0dfb87e5058276ab515b626e0c72c3cfa075189f Mon Sep 17 00:00:00 2001 From: Lazizbek Ergashev Date: Sat, 7 Nov 2026 04:17:30 +0000 Subject: [PATCH] fix(api): keep segment attachment files in clear-orphaned-file-records (#45750) --- api/commands/storage.py | 2 + api/tests/unit_tests/commands/test_storage.py | 49 +++++++++++++++++++ 2 files changed, 51 insertions(+) create mode 100644 api/tests/unit_tests/commands/test_storage.py diff --git a/api/commands/storage.py b/api/commands/storage.py index a21b66aa4e8..9f95485e3e3 100644 --- a/api/commands/storage.py +++ b/api/commands/storage.py @@ -28,6 +28,7 @@ def clear_orphaned_file_records(force: bool): ] ids_tables = [ {"type": "uuid", "table": "message_files", "column": "upload_file_id"}, + {"type": "uuid", "table": "segment_attachment_bindings", "column": "attachment_id"}, {"type": "text", "table": "documents", "column": "data_source_info"}, {"type": "text", "table": "document_segments", "column": "content"}, {"type": "text", "table": "messages", "column": "answer"}, @@ -378,6 +379,7 @@ def file_usage( ] ids_tables = [ {"type": "uuid", "table": "message_files", "column": "upload_file_id", "pk_column": "id"}, + {"type": "uuid", "table": "segment_attachment_bindings", "column": "attachment_id", "pk_column": "id"}, {"type": "text", "table": "documents", "column": "data_source_info", "pk_column": "id"}, {"type": "text", "table": "document_segments", "column": "content", "pk_column": "id"}, {"type": "text", "table": "messages", "column": "answer", "pk_column": "id"}, diff --git a/api/tests/unit_tests/commands/test_storage.py b/api/tests/unit_tests/commands/test_storage.py new file mode 100644 index 00000000000..01c9bfb3845 --- /dev/null +++ b/api/tests/unit_tests/commands/test_storage.py @@ -0,0 +1,49 @@ +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass + +import pytest +from click.testing import CliRunner +from sqlalchemy import TextClause + +from commands import storage + + +class _FakeEngine: + def __init__(self, rows_by_table: dict[str, list[tuple[str, ...]]]) -> None: + self.rows_by_table = rows_by_table + self.executed: list[str] = [] + + @contextmanager + def begin(self) -> Generator["_FakeEngine"]: + yield self + + def execute(self, statement: TextClause, *_: object) -> list[tuple[str, ...]]: + query = str(statement) + self.executed.append(query) + if query.startswith("SELECT"): + table = query.split(" FROM ")[1].split()[0] + return self.rows_by_table.get(table, []) + return [] + + +@dataclass +class _FakeDb: + engine: _FakeEngine + + +def test_clear_orphaned_file_records_keeps_files_used_as_segment_attachments(monkeypatch: pytest.MonkeyPatch) -> None: + attachment_id = "0198a6b2-6e4c-7c3a-9d1e-2f3b4c5d6e7f" + engine = _FakeEngine( + { + "upload_files": [(attachment_id, "upload_files/tenant/image.png")], + "segment_attachment_bindings": [(attachment_id,)], + } + ) + monkeypatch.setattr(storage, "db", _FakeDb(engine)) + + result = CliRunner().invoke(storage.clear_orphaned_file_records, ["--force"]) + + assert result.exit_code == 0 + assert "No orphaned file records found" in result.output + assert not any(query.startswith("DELETE FROM upload_files") for query in engine.executed)