test: use sqlite3 session in test_remote_fetcher (#38773)

api/tests/unit_tests/core/file/test_remote_fetcher.py

@@ -2,20 +2,104 @@ import base64
 import hashlib
 import hmac
 import urllib.parse
-from types import SimpleNamespace
+from collections.abc import Iterator
+from dataclasses import dataclass
+from datetime import UTC, datetime
 from unittest.mock import MagicMock
 
 import httpx
 import pytest
+from sqlalchemy import Engine, event
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.datasource.datasource_file_manager import DatasourceFileManager
 from core.file import remote_fetcher
 from core.tools.signature import sign_tool_file, sign_upload_file_preview_url
 from core.tools.tool_file_manager import ToolFileManager
+from extensions.storage.storage_type import StorageType
+from models import CreatorUserRole, ToolFile, UploadFile
 
 UPLOAD_FILE_ID = "1602650a-4fe4-423c-85a2-af76c083e3c4"
 TOOL_FILE_ID = "2602650a-4fe4-423c-85a2-af76c083e3c4"
 DATASOURCE_FILE_ID = "3602650a-4fe4-423c-85a2-af76c083e3c4"
+TENANT_ID = "tenant-1"
+USER_ID = "user-1"
+
+
+@dataclass(frozen=True)
+class FileDatabase:
+    session: Session
+    upload_file: UploadFile
+    tool_file: ToolFile
+    datasource_upload_file: UploadFile
+    statements: list[str]
+
+
+@pytest.fixture(autouse=True)
+def file_database(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> Iterator[FileDatabase]:
+    """Persist file records and bind remote-fetcher sessions to SQLite."""
+    UploadFile.metadata.create_all(sqlite_engine, tables=[UploadFile.__table__, ToolFile.__table__])
+    sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
+    monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker)
+
+    upload_file = UploadFile(
+        tenant_id=TENANT_ID,
+        storage_type=StorageType.LOCAL,
+        key="upload_files/tenant/hello.txt",
+        name="hello.txt",
+        size=5,
+        extension="txt",
+        mime_type="text/plain",
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by=USER_ID,
+        created_at=datetime(2026, 1, 1, tzinfo=UTC),
+        used=False,
+    )
+    upload_file.id = UPLOAD_FILE_ID
+    tool_file = ToolFile(
+        user_id=USER_ID,
+        tenant_id=TENANT_ID,
+        conversation_id=None,
+        file_key="tools/tenant/result.txt",
+        mimetype="text/plain",
+        name="result.txt",
+        size=6,
+    )
+    tool_file.id = TOOL_FILE_ID
+    datasource_upload_file = UploadFile(
+        tenant_id=TENANT_ID,
+        storage_type=StorageType.LOCAL,
+        key="datasources/tenant/data.txt",
+        name="data.txt",
+        size=4,
+        extension="txt",
+        mime_type="text/plain",
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by=USER_ID,
+        created_at=datetime(2026, 1, 1, tzinfo=UTC),
+        used=False,
+    )
+    datasource_upload_file.id = DATASOURCE_FILE_ID
+
+    statements: list[str] = []
+
+    def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None:
+        statements.append(statement)
+
+    with sqlite_session_maker() as session:
+        session.add_all([upload_file, tool_file, datasource_upload_file])
+        session.commit()
+        event.listen(sqlite_engine, "before_cursor_execute", record_statement)
+        try:
+            yield FileDatabase(
+                session=session,
+                upload_file=upload_file,
+                tool_file=tool_file,
+                datasource_upload_file=datasource_upload_file,
+                statements=statements,
+            )
+        finally:
+            event.remove(sqlite_engine, "before_cursor_execute", record_statement)
 
 
 def _signed_url(*, base_url: str, path: str, payload: str, secret: str = "test-secret") -> str:
@@ -44,15 +128,6 @@ def _patch_file_fetcher_config(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(remote_fetcher.time, "time", lambda: 1700000100)
 
 
-def _patch_session(monkeypatch: pytest.MonkeyPatch):
-    session = MagicMock()
-    session_cm = MagicMock()
-    session_cm.__enter__.return_value = session
-    session_cm.__exit__.return_value = False
-    monkeypatch.setattr(remote_fetcher.session_factory, "create_session", MagicMock(return_value=session_cm))
-    return session
-
-
 def _patch_ssrf_make_request(monkeypatch: pytest.MonkeyPatch, response=None):
     make_request = MagicMock(return_value=response) if response is not None else MagicMock()
     monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", make_request)
@@ -67,18 +142,7 @@ def _patch_signer_times(monkeypatch: pytest.MonkeyPatch):
 
 def test_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    session = _patch_session(monkeypatch)
-    upload_file = SimpleNamespace(
-        id=UPLOAD_FILE_ID,
-        key="upload_files/tenant/hello.txt",
-        name="hello.txt",
-        mime_type="text/plain",
-        size=5,
-        extension="txt",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
     load_once = MagicMock(return_value=b"hello")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -94,29 +158,22 @@ def test_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pyte
     assert response.headers["Content-Type"] == "text/plain"
     assert response.headers["Content-Length"] == "5"
     assert response.request.method == "GET"
-    get_upload_file.assert_called_once_with(
-        session=session,
-        file_id=UPLOAD_FILE_ID,
-    )
     load_once.assert_called_once_with("upload_files/tenant/hello.txt")
     ssrf_make_request.assert_not_called()
 
 
-def test_make_request_resolves_upload_preview_url_generated_by_signer(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_resolves_upload_preview_url_generated_by_signer(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
     _patch_signer_times(monkeypatch)
-    session = _patch_session(monkeypatch)
+    file_database.upload_file.key = "upload_files/tenant/image.png"
-    upload_file = SimpleNamespace(
+    file_database.upload_file.name = "image.png"
-        id=UPLOAD_FILE_ID,
+    file_database.upload_file.mime_type = "image/png"
-        key="upload_files/tenant/image.png",
+    file_database.upload_file.size = 6
-        name="image.png",
+    file_database.upload_file.extension = "png"
-        mime_type="image/png",
+    file_database.session.commit()
-        size=6,
-        extension=".png",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
     load_once = MagicMock(return_value=b"image!")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = MagicMock()
     monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request)
@@ -127,25 +184,21 @@ def test_make_request_resolves_upload_preview_url_generated_by_signer(monkeypatc
     assert response.status_code == 200
     assert response.content == b"image!"
     assert response.headers["Content-Type"] == "image/png"
-    get_upload_file.assert_called_once_with(session=session, file_id=UPLOAD_FILE_ID)
     load_once.assert_called_once_with("upload_files/tenant/image.png")
     ssrf_make_request.assert_not_called()
 
 
-def test_make_request_resolves_sign_tool_file_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_resolves_sign_tool_file_url_with_empty_extension(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
     _patch_signer_times(monkeypatch)
-    session = _patch_session(monkeypatch)
+    file_database.tool_file.file_key = "tools/tenant/no-extension"
-    tool_file = SimpleNamespace(
+    file_database.tool_file.name = "no-extension"
-        id=TOOL_FILE_ID,
+    file_database.tool_file.mimetype = "application/octet-stream"
-        file_key="tools/tenant/no-extension",
+    file_database.tool_file.size = 8
-        name="no-extension",
+    file_database.session.commit()
-        mimetype="application/octet-stream",
-        size=8,
-    )
-    get_tool_file = MagicMock(return_value=tool_file)
     load_once = MagicMock(return_value=b"tooldata")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = MagicMock()
     monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request)
@@ -156,25 +209,21 @@ def test_make_request_resolves_sign_tool_file_url_with_empty_extension(monkeypat
     assert response.status_code == 200
     assert response.content == b"tooldata"
     assert response.headers["Content-Type"] == "application/octet-stream"
-    get_tool_file.assert_called_once_with(session=session, file_id=TOOL_FILE_ID)
     load_once.assert_called_once_with("tools/tenant/no-extension")
     ssrf_make_request.assert_not_called()
 
 
-def test_make_request_resolves_tool_manager_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_resolves_tool_manager_url_with_empty_extension(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
     _patch_signer_times(monkeypatch)
-    session = _patch_session(monkeypatch)
+    file_database.tool_file.file_key = "tools/tenant/manager-file"
-    tool_file = SimpleNamespace(
+    file_database.tool_file.name = "manager-file"
-        id=TOOL_FILE_ID,
+    file_database.tool_file.mimetype = "application/octet-stream"
-        file_key="tools/tenant/manager-file",
+    file_database.tool_file.size = 12
-        name="manager-file",
+    file_database.session.commit()
-        mimetype="application/octet-stream",
-        size=12,
-    )
-    get_tool_file = MagicMock(return_value=tool_file)
     load_once = MagicMock(return_value=b"manager-data")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = MagicMock()
     monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request)
@@ -184,28 +233,22 @@ def test_make_request_resolves_tool_manager_url_with_empty_extension(monkeypatch
 
     assert response.status_code == 200
     assert response.content == b"manager-data"
-    get_tool_file.assert_called_once_with(session=session, file_id=TOOL_FILE_ID)
     load_once.assert_called_once_with("tools/tenant/manager-file")
     ssrf_make_request.assert_not_called()
 
 
-def test_make_request_resolves_datasource_manager_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_resolves_datasource_manager_url_with_empty_extension(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
     _patch_signer_times(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.datasource_upload_file.key = "datasources/tenant/no-extension"
-    upload_file = SimpleNamespace(
+    file_database.datasource_upload_file.name = "no-extension"
-        id=DATASOURCE_FILE_ID,
+    file_database.datasource_upload_file.mime_type = "application/octet-stream"
-        key="datasources/tenant/no-extension",
+    file_database.datasource_upload_file.size = 10
-        name="no-extension",
+    file_database.datasource_upload_file.extension = ""
-        mime_type="application/octet-stream",
+    file_database.session.commit()
-        size=10,
-        extension="",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
-    get_tool_file = MagicMock()
     load_once = MagicMock(return_value=b"datasource")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = MagicMock()
     monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request)
@@ -216,25 +259,12 @@ def test_make_request_resolves_datasource_manager_url_with_empty_extension(monke
     assert response.status_code == 200
     assert response.content == b"datasource"
     assert response.headers["Content-Type"] == "application/octet-stream"
-    get_upload_file.assert_called_once()
-    get_tool_file.assert_not_called()
     load_once.assert_called_once_with("datasources/tenant/no-extension")
     ssrf_make_request.assert_not_called()
 
 
 def test_head_signed_upload_file_url_returns_metadata_without_storage_content(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    session = _patch_session(monkeypatch)
-    upload_file = SimpleNamespace(
-        id=UPLOAD_FILE_ID,
-        key="upload_files/tenant/hello.txt",
-        name="hello.txt",
-        mime_type="text/plain",
-        size=5,
-        extension="txt",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     load_once = MagicMock(return_value=b"hello")
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
@@ -251,28 +281,13 @@ def test_head_signed_upload_file_url_returns_metadata_without_storage_content(mo
     assert response.headers["Content-Type"] == "text/plain"
     assert response.headers["Content-Length"] == "5"
     assert response.request.method == "HEAD"
-    get_upload_file.assert_called_once_with(
-        session=session,
-        file_id=UPLOAD_FILE_ID,
-    )
     load_once.assert_not_called()
     ssrf_make_request.assert_not_called()
 
 
 def test_make_request_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
-    upload_file = SimpleNamespace(
-        id=UPLOAD_FILE_ID,
-        key="upload_files/tenant/hello.txt",
-        name="hello.txt",
-        mime_type="text/plain",
-        size=5,
-        extension="txt",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
     load_once = MagicMock(return_value=b"hello")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -286,24 +301,12 @@ def test_make_request_get_signed_upload_file_url_reads_storage_without_ssrf(monk
     assert response.status_code == 200
     assert response.content == b"hello"
     assert response.request.method == "GET"
-    get_upload_file.assert_called_once()
     load_once.assert_called_once_with("upload_files/tenant/hello.txt")
     ssrf_make_request.assert_not_called()
 
 
 def test_make_request_head_signed_upload_file_url_returns_metadata_without_ssrf(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
-    upload_file = SimpleNamespace(
-        id=UPLOAD_FILE_ID,
-        key="upload_files/tenant/hello.txt",
-        name="hello.txt",
-        mime_type="text/plain",
-        size=5,
-        extension="txt",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     load_once = MagicMock(return_value=b"hello")
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
@@ -320,15 +323,14 @@ def test_make_request_head_signed_upload_file_url_returns_metadata_without_ssrf(
     assert response.headers["Content-Type"] == "text/plain"
     assert response.headers["Content-Length"] == "5"
     assert response.request.method == "HEAD"
-    get_upload_file.assert_called_once()
     load_once.assert_not_called()
     ssrf_make_request.assert_not_called()
 
 
-def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    get_upload_file = MagicMock()
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     url = f"http://localhost:5001/files/{UPLOAD_FILE_ID}/file-preview?timestamp=1700000000&nonce=nonce"
     proxy_response = httpx.Response(403, request=httpx.Request("GET", url))
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch, proxy_response)
@@ -336,7 +338,7 @@ def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch:
     response = remote_fetcher.make_request("GET", url, timeout=3)
 
     assert response is proxy_response
-    get_upload_file.assert_not_called()
+    assert file_database.statements == []
     ssrf_make_request.assert_called_once_with(
         method="GET",
         url=url,
@@ -345,10 +347,10 @@ def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch:
     )
 
 
-def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkeypatch: pytest.MonkeyPatch):
+def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    get_upload_file = MagicMock()
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     proxy_response = httpx.Response(201, request=httpx.Request("POST", f"http://localhost:5001/files/{UPLOAD_FILE_ID}"))
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch, proxy_response)
     url = _signed_url(
@@ -360,7 +362,7 @@ def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkey
     response = remote_fetcher.make_request("POST", url, json={"name": "ignored"})
 
     assert response is proxy_response
-    get_upload_file.assert_not_called()
+    assert file_database.statements == []
     ssrf_make_request.assert_called_once_with(
         method="POST",
         url=url,
@@ -369,20 +371,17 @@ def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkey
     )
 
 
-def test_get_signed_image_preview_url_uses_image_preview_signature(monkeypatch: pytest.MonkeyPatch):
+def test_get_signed_image_preview_url_uses_image_preview_signature(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.upload_file.key = "upload_files/tenant/image.png"
-    upload_file = SimpleNamespace(
+    file_database.upload_file.name = "image.png"
-        id=UPLOAD_FILE_ID,
+    file_database.upload_file.mime_type = "image/png"
-        key="upload_files/tenant/image.png",
+    file_database.upload_file.size = 6
-        name="image.png",
+    file_database.upload_file.extension = "png"
-        mime_type="image/png",
+    file_database.session.commit()
-        size=6,
-        extension="png",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
     load_once = MagicMock(return_value=b"image!")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -396,7 +395,6 @@ def test_get_signed_image_preview_url_uses_image_preview_signature(monkeypatch:
     assert response.status_code == 200
     assert response.content == b"image!"
     assert response.headers["Content-Type"] == "image/png"
-    get_upload_file.assert_called_once()
     load_once.assert_called_once_with("upload_files/tenant/image.png")
     ssrf_make_request.assert_not_called()
 
@@ -597,11 +595,12 @@ def test_invalid_configured_file_origin_delegates_to_ssrf_proxy(monkeypatch: pyt
     )
 
 
-def test_signed_upload_file_url_returns_404_when_record_missing(monkeypatch: pytest.MonkeyPatch):
+def test_signed_upload_file_url_returns_404_when_record_missing(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.session.delete(file_database.upload_file)
-    get_upload_file = MagicMock(return_value=None)
+    file_database.session.commit()
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
         base_url="http://localhost:5001",
@@ -613,23 +612,12 @@ def test_signed_upload_file_url_returns_404_when_record_missing(monkeypatch: pyt
 
     assert response.status_code == 404
     assert response.content == b""
-    get_upload_file.assert_called_once()
     ssrf_make_request.assert_not_called()
 
 
 def test_get_signed_tool_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    session = _patch_session(monkeypatch)
-    tool_file = SimpleNamespace(
-        id=TOOL_FILE_ID,
-        file_key="tools/tenant/result.txt",
-        name="result.txt",
-        mimetype="text/plain",
-        size=6,
-    )
-    get_tool_file = MagicMock(return_value=tool_file)
     load_once = MagicMock(return_value=b"result")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -643,19 +631,16 @@ def test_get_signed_tool_file_url_reads_storage_without_ssrf(monkeypatch: pytest
     assert response.status_code == 200
     assert response.content == b"result"
     assert response.headers["Content-Type"] == "text/plain"
-    get_tool_file.assert_called_once_with(
-        session=session,
-        file_id=TOOL_FILE_ID,
-    )
     load_once.assert_called_once_with("tools/tenant/result.txt")
     ssrf_make_request.assert_not_called()
 
 
-def test_signed_tool_file_url_returns_404_when_record_missing(monkeypatch: pytest.MonkeyPatch):
+def test_signed_tool_file_url_returns_404_when_record_missing(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.session.delete(file_database.tool_file)
-    get_tool_file = MagicMock(return_value=None)
+    file_database.session.commit()
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
         base_url="http://localhost:5001",
@@ -667,26 +652,12 @@ def test_signed_tool_file_url_returns_404_when_record_missing(monkeypatch: pytes
 
     assert response.status_code == 404
     assert response.content == b""
-    get_tool_file.assert_called_once()
     ssrf_make_request.assert_not_called()
 
 
 def test_get_signed_datasource_file_url_reads_upload_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
-    upload_file = SimpleNamespace(
-        id=DATASOURCE_FILE_ID,
-        key="datasources/tenant/data.txt",
-        name="data.txt",
-        mime_type="text/plain",
-        size=4,
-        extension="txt",
-    )
-    get_upload_file = MagicMock(return_value=upload_file)
-    get_tool_file = MagicMock(return_value=None)
     load_once = MagicMock(return_value=b"data")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -699,27 +670,28 @@ def test_get_signed_datasource_file_url_reads_upload_storage_without_ssrf(monkey
 
     assert response.status_code == 200
     assert response.content == b"data"
-    get_upload_file.assert_called_once()
-    get_tool_file.assert_not_called()
     load_once.assert_called_once_with("datasources/tenant/data.txt")
     ssrf_make_request.assert_not_called()
 
 
-def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing(monkeypatch: pytest.MonkeyPatch):
+def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.session.delete(file_database.datasource_upload_file)
-    tool_file = SimpleNamespace(
+    datasource_tool_file = ToolFile(
-        id=DATASOURCE_FILE_ID,
+        user_id=USER_ID,
+        tenant_id=TENANT_ID,
+        conversation_id=None,
         file_key="datasources/tenant/tool-data.txt",
         name="tool-data.txt",
         mimetype="text/plain",
         size=9,
     )
-    get_upload_file = MagicMock(return_value=None)
+    datasource_tool_file.id = DATASOURCE_FILE_ID
-    get_tool_file = MagicMock(return_value=tool_file)
+    file_database.session.add(datasource_tool_file)
+    file_database.session.commit()
     load_once = MagicMock(return_value=b"tool-data")
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
@@ -734,19 +706,16 @@ def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing(m
     assert response.content == b"tool-data"
     assert response.headers["Content-Type"] == "text/plain"
     assert response.headers["Content-Length"] == "9"
-    get_upload_file.assert_called_once()
-    get_tool_file.assert_called_once()
     load_once.assert_called_once_with("datasources/tenant/tool-data.txt")
     ssrf_make_request.assert_not_called()
 
 
-def test_signed_datasource_file_url_returns_404_when_records_missing(monkeypatch: pytest.MonkeyPatch):
+def test_signed_datasource_file_url_returns_404_when_records_missing(
+    monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase
+):
     _patch_file_fetcher_config(monkeypatch)
-    _patch_session(monkeypatch)
+    file_database.session.delete(file_database.datasource_upload_file)
-    get_upload_file = MagicMock(return_value=None)
+    file_database.session.commit()
-    get_tool_file = MagicMock(return_value=None)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file)
-    monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file)
     ssrf_make_request = _patch_ssrf_make_request(monkeypatch)
     url = _signed_url(
         base_url="http://localhost:5001",
@@ -758,8 +727,6 @@ def test_signed_datasource_file_url_returns_404_when_records_missing(monkeypatch
 
     assert response.status_code == 404
     assert response.content == b""
-    get_upload_file.assert_called_once()
-    get_tool_file.assert_called_once()
     ssrf_make_request.assert_not_called()