From 8ad131bb3badc8f66b687b4d629ffaf14fb505b4 Mon Sep 17 00:00:00 2001 From: aliworksx08 <57456290+aliworksx08@users.noreply.github.com> Date: Thu, 9 Apr 2026 09:15:59 -0500 Subject: [PATCH] refactor: migrate session.query to select API in file service (#34852) --- api/services/file_service.py | 14 ++++----- .../unit_tests/services/test_file_service.py | 30 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/api/services/file_service.py b/api/services/file_service.py index 50a326d813..7443ca3271 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -132,8 +132,8 @@ class FileService: return file_size <= file_size_limit def get_file_base64(self, file_id: str) -> str: - upload_file = ( - self._session_maker(expire_on_commit=False).query(UploadFile).where(UploadFile.id == file_id).first() + upload_file = self._session_maker(expire_on_commit=False).scalar( + select(UploadFile).where(UploadFile.id == file_id).limit(1) ) if not upload_file: raise NotFound("File not found") @@ -178,7 +178,7 @@ class FileService: Return a short text preview extracted from a document file. """ with self._session_maker(expire_on_commit=False) as session: - upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first() + upload_file = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if not upload_file: raise NotFound("File not found") @@ -200,7 +200,7 @@ class FileService: if not result: raise NotFound("File not found or signature is invalid") with self._session_maker(expire_on_commit=False) as session: - upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first() + upload_file = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if not upload_file: raise NotFound("File not found or signature is invalid") @@ -220,7 +220,7 @@ class FileService: raise NotFound("File not found or signature is invalid") with self._session_maker(expire_on_commit=False) as session: - upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first() + upload_file = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if not upload_file: raise NotFound("File not found or signature is invalid") @@ -231,7 +231,7 @@ class FileService: def get_public_image_preview(self, file_id: str): with self._session_maker(expire_on_commit=False) as session: - upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first() + upload_file = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if not upload_file: raise NotFound("File not found or signature is invalid") @@ -247,7 +247,7 @@ class FileService: def get_file_content(self, file_id: str) -> str: with self._session_maker(expire_on_commit=False) as session: - upload_file: UploadFile | None = session.query(UploadFile).where(UploadFile.id == file_id).first() + upload_file: UploadFile | None = session.scalar(select(UploadFile).where(UploadFile.id == file_id).limit(1)) if not upload_file: raise NotFound("File not found") diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index b7259c3e82..8e1b22886b 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -165,7 +165,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "test_key" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load_once.return_value = b"test content" @@ -178,7 +178,7 @@ class TestFileService: mock_storage.load_once.assert_called_once_with("test_key") def test_get_file_base64_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_base64("non_existent") @@ -215,7 +215,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "pdf" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with patch("services.file_service.ExtractProcessor.load_from_upload_file") as mock_extract: mock_extract.return_value = "Extracted text content" @@ -227,7 +227,7 @@ class TestFileService: assert result == "Extracted text content" def test_get_file_preview_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_preview("non_existent") @@ -235,7 +235,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "exe" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with pytest.raises(UnsupportedFileTypeError): file_service.get_file_preview("file_id") @@ -246,7 +246,7 @@ class TestFileService: upload_file.extension = "jpg" upload_file.mime_type = "image/jpeg" upload_file.key = "key" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with ( patch("services.file_service.file_helpers.verify_image_signature") as mock_verify, @@ -269,7 +269,7 @@ class TestFileService: file_service.get_image_preview("file_id", "ts", "nonce", "sign") def test_get_image_preview_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(NotFound, match="File not found or signature is invalid"): @@ -279,7 +279,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "txt" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with patch("services.file_service.file_helpers.verify_image_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(UnsupportedFileTypeError): @@ -289,7 +289,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "key" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with ( patch("services.file_service.file_helpers.verify_file_signature") as mock_verify, @@ -309,7 +309,7 @@ class TestFileService: file_service.get_file_generator_by_file_id("file_id", "ts", "nonce", "sign") def test_get_file_generator_by_file_id_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with patch("services.file_service.file_helpers.verify_file_signature") as mock_verify: mock_verify.return_value = True with pytest.raises(NotFound, match="File not found or signature is invalid"): @@ -321,7 +321,7 @@ class TestFileService: upload_file.extension = "png" upload_file.mime_type = "image/png" upload_file.key = "key" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load.return_value = b"image content" @@ -330,7 +330,7 @@ class TestFileService: assert mime == "image/png" def test_get_public_image_preview_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found or signature is invalid"): file_service.get_public_image_preview("file_id") @@ -338,7 +338,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.extension = "txt" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with pytest.raises(UnsupportedFileTypeError): file_service.get_public_image_preview("file_id") @@ -346,7 +346,7 @@ class TestFileService: upload_file = MagicMock(spec=UploadFile) upload_file.id = "file_id" upload_file.key = "key" - mock_db_session.query().where().first.return_value = upload_file + mock_db_session.scalar.return_value = upload_file with patch("services.file_service.storage") as mock_storage: mock_storage.load.return_value = b"hello world" @@ -354,7 +354,7 @@ class TestFileService: assert result == "hello world" def test_get_file_content_not_found(self, file_service, mock_db_session): - mock_db_session.query().where().first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(NotFound, match="File not found"): file_service.get_file_content("file_id")