opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-04-15 18:06:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: migrate session.query to select API in file service (#34852)

Browse Source
This commit is contained in:
aliworksx08 2026-04-09 09:15:59 -05:00 committed by GitHub
parent 41eeb1f2e7
commit 8ad131bb3b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
api/services/file_service.py
View File

@ -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")

30
api/tests/unit_tests/services/test_file_service.py
View File

@ -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")