diff --git a/api/controllers/console/files.py b/api/controllers/console/files.py index 2f166340846..9ae0bafb6c7 100644 --- a/api/controllers/console/files.py +++ b/api/controllers/console/files.py @@ -17,26 +17,30 @@ from controllers.common.errors import ( UnsupportedFileTypeError, ) from controllers.common.fields import AllowedExtensionsResponse, TextContentResponse -from controllers.common.schema import register_response_schema_models, register_schema_models +from controllers.common.schema import JsonResponseWithStatus, register_response_schema_models +from controllers.console import console_ns +from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, setup_required, - with_current_tenant_id, with_current_user, ) -from extensions.ext_database import db +from extensions.ext_application_services import application_services from fields.file_fields import FileResponse, UploadConfig from libs.helper import dump_response from libs.login import login_required +from machinery.context import RequestContext from models import Account, UploadFile from services.feature_service import FeatureService -from services.file_service import FileService -from . import console_ns - -register_schema_models(console_ns, UploadConfig, FileResponse) -register_response_schema_models(console_ns, AllowedExtensionsResponse, TextContentResponse) +register_response_schema_models( + console_ns, + UploadConfig, + FileResponse, + AllowedExtensionsResponse, + TextContentResponse, +) PREVIEW_WORDS_LIMIT = 3000 @@ -70,7 +74,7 @@ def upload_file_from_request(*, current_user: Account, resource_tenant_id: str | file = request.files["file"] if not file.filename: - raise FilenameNotExistsError + raise FilenameNotExistsError() if source == "datasets" and not current_user.is_dataset_editor: raise Forbidden() @@ -84,7 +88,7 @@ def upload_file_from_request(*, current_user: Account, resource_tenant_id: str | ) try: - return FileService(db.engine).upload_file( + return application_services().files.upload_file( filename=file.filename, content=file.stream.read(), mimetype=file.mimetype, @@ -94,24 +98,21 @@ def upload_file_from_request(*, current_user: Account, resource_tenant_id: str | default_file_size_limit=default_file_size_limit, ) except services.errors.file.FileTooLargeError as file_too_large_error: - raise FileTooLargeError(file_too_large_error.description) - except services.errors.file.UnsupportedFileTypeError: - raise UnsupportedFileTypeError() + raise FileTooLargeError(file_too_large_error.description) from file_too_large_error + except services.errors.file.UnsupportedFileTypeError as unsupported_file_type_error: + raise UnsupportedFileTypeError() from unsupported_file_type_error except services.errors.file.BlockedFileExtensionError as blocked_extension_error: - raise BlockedFileExtensionError(blocked_extension_error.description) + raise BlockedFileExtensionError(blocked_extension_error.description) from blocked_extension_error @console_ns.route("/files/upload") class FileApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[UploadConfig.__name__]) - @with_current_tenant_id - def get(self, current_tenant_id: str): + @console_account_admission() + def get(self, request_context: RequestContext) -> JsonResponseWithStatus: config = UploadConfig( file_size_limit=dify_config.UPLOAD_FILE_SIZE_LIMIT, - knowledge_file_size_limit=FeatureService.get_knowledge_file_size_limit(current_tenant_id), + knowledge_file_size_limit=FeatureService.get_knowledge_file_size_limit(request_context.active_workspace_id), batch_count_limit=dify_config.UPLOAD_FILE_BATCH_LIMIT, file_upload_limit=dify_config.BATCH_UPLOAD_LIMIT, image_file_size_limit=dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT, @@ -123,7 +124,7 @@ class FileApi(Resource): single_chunk_attachment_limit=dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT, attachment_image_file_size_limit=dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT, ) - return config.model_dump(mode="json"), 200 + return dump_response(UploadConfig, config), 200 @setup_required @login_required @@ -132,7 +133,7 @@ class FileApi(Resource): @console_ns.doc(consumes=["multipart/form-data"], params=FILE_UPLOAD_PARAMS) @console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__]) @with_current_user - def post(self, current_user: Account): + def post(self, current_user: Account) -> JsonResponseWithStatus: upload_file = upload_file_from_request(current_user=current_user) return dump_response(FileResponse, upload_file), 201 @@ -140,22 +141,21 @@ class FileApi(Resource): @console_ns.route("/files//preview") class FilePreviewApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__]) - @with_current_tenant_id - def get(self, current_tenant_id: str, file_id: UUID): + @console_account_admission() + def get(self, request_context: RequestContext, file_id: UUID) -> dict[str, object]: + current_tenant_id = request_context.active_workspace_id file_id_str = str(file_id) - text = FileService(db.engine).get_file_preview(file_id_str, current_tenant_id) - return TextContentResponse(content=text).model_dump(mode="json") + text = application_services().files.get_file_preview(file_id=file_id_str, tenant_id=current_tenant_id) + return dump_response(TextContentResponse, {"content": text}) @console_ns.route("/files/support-type") class FileSupportTypeApi(Resource): - @setup_required - @login_required - @account_initialization_required @console_ns.response(200, "Success", console_ns.models[AllowedExtensionsResponse.__name__]) - def get(self): - return AllowedExtensionsResponse(allowed_extensions=list(DOCUMENT_EXTENSIONS)).model_dump(mode="json") + @console_account_admission() + def get(self, _request_context: RequestContext) -> dict[str, object]: + return dump_response( + AllowedExtensionsResponse, + {"allowed_extensions": list(DOCUMENT_EXTENSIONS)}, + ) diff --git a/api/tests/unit_tests/controllers/console/test_files.py b/api/tests/unit_tests/controllers/console/test_files.py index 60bfb92d1c1..e5db1174b83 100644 --- a/api/tests/unit_tests/controllers/console/test_files.py +++ b/api/tests/unit_tests/controllers/console/test_files.py @@ -4,7 +4,6 @@ from unittest.mock import patch import pytest from flask import Flask -from sqlalchemy import Engine from werkzeug.exceptions import Forbidden from configs import dify_config @@ -24,6 +23,7 @@ from controllers.console.files import ( upload_file_from_request, ) from extensions.storage.storage_type import StorageType +from machinery.context import RequestContext from models import Account from models.account import AccountStatus, TenantAccountRole from models.enums import CreatorUserRole @@ -57,6 +57,15 @@ def _upload_file(*, file_id: str = "file-id-123", size: int = 1024) -> UploadFil return upload_file +def _request_context(*, workspace_id: str = "tenant-1") -> RequestContext: + return RequestContext( + request_id="request-1", + trace_id="trace-1", + account_id="account-1", + active_workspace_id=workspace_id, + ) + + @pytest.fixture def app(): app = Flask(__name__) @@ -92,17 +101,9 @@ def mock_account_context(mock_current_user): @pytest.fixture -def mock_db(sqlite_engine: Engine): - with patch("controllers.console.files.db") as db_mock: - db_mock.engine = sqlite_engine - yield db_mock - - -@pytest.fixture -def mock_file_service(mock_db): - with patch("controllers.console.files.FileService") as fs: - instance = fs.return_value - yield instance +def mock_file_service(): + with patch("controllers.console.files.application_services") as services: + yield services.return_value.files class TestFileApiGet: @@ -117,7 +118,7 @@ class TestFileApiGet: return_value=50, ) as get_knowledge_file_size_limit, ): - data, status = get_method(api, "tenant-1") + data, status = get_method(api, _request_context()) assert status == 200 assert "file_size_limit" in data @@ -196,6 +197,15 @@ class TestFileApiPost: assert status == 201 assert response["id"] == "file-id-123" assert response["name"] == "test.txt" + mock_file_service.upload_file.assert_called_once_with( + filename="test.txt", + content=b"hello", + mimetype="text/plain", + user=mock_account_context, + tenant_id=None, + source=None, + default_file_size_limit=None, + ) def test_upload_with_resource_tenant(self, app: Flask, mock_account_context, mock_file_service): upload_file = _upload_file() @@ -276,9 +286,11 @@ class TestFileApiPost: } with app.test_request_context(method="POST", data=data): - with pytest.raises(FileTooLargeError): + with pytest.raises(FileTooLargeError) as error_info: post_method(api, mock_account_context) + assert error_info.value.__cause__ is error + def test_unsupported_file_type(self, app: Flask, mock_account_context, mock_file_service): api = FileApi() post_method = unwrap(api.post) @@ -293,9 +305,11 @@ class TestFileApiPost: } with app.test_request_context(method="POST", data=data): - with pytest.raises(UnsupportedFileTypeError): + with pytest.raises(UnsupportedFileTypeError) as error_info: post_method(api, mock_account_context) + assert error_info.value.__cause__ is error + def test_blocked_extension(self, app: Flask, mock_account_context, mock_file_service): api = FileApi() post_method = unwrap(api.post) @@ -310,9 +324,12 @@ class TestFileApiPost: } with app.test_request_context(method="POST", data=data): - with pytest.raises(BlockedFileExtensionError): + with pytest.raises(BlockedFileExtensionError) as error_info: post_method(api, mock_account_context) + assert error_info.value.description == error.description + assert error_info.value.__cause__ is error + class TestFilePreviewApi: def test_get_preview(self, app: Flask, mock_account_context, mock_file_service): @@ -321,9 +338,10 @@ class TestFilePreviewApi: mock_file_service.get_file_preview.return_value = "preview text" with app.test_request_context(): - result = get_method(api, "tenant-123", "1234") + result = get_method(api, _request_context(workspace_id="tenant-123"), "1234") assert result == {"content": "preview text"} + mock_file_service.get_file_preview.assert_called_once_with(file_id="1234", tenant_id="tenant-123") class TestFileSupportTypeApi: @@ -332,6 +350,6 @@ class TestFileSupportTypeApi: get_method = unwrap(api.get) with app.test_request_context(): - result = get_method(api) + result = get_method(api, _request_context()) assert result == {"allowed_extensions": list(DOCUMENT_EXTENSIONS)} diff --git a/api/tests/unit_tests/controllers/console/test_files_security.py b/api/tests/unit_tests/controllers/console/test_files_security.py index eb090ff8ea5..7bab15a5a95 100644 --- a/api/tests/unit_tests/controllers/console/test_files_security.py +++ b/api/tests/unit_tests/controllers/console/test_files_security.py @@ -1,6 +1,5 @@ import builtins import io -from unittest.mock import patch import pytest from flask.views import MethodView @@ -8,15 +7,11 @@ from werkzeug.exceptions import Forbidden from controllers.common.errors import ( FilenameNotExistsError, - FileTooLargeError, NoFileUploadedError, TooManyFilesError, - UnsupportedFileTypeError, ) from models import Account from models.account import AccountStatus, TenantAccountRole -from services.errors.file import FileTooLargeError as ServiceFileTooLargeError -from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -128,34 +123,7 @@ class TestFileUploadSecurity: raise Forbidden() # Test passes if no exception is raised - # Test 4: Service error handling - @patch("controllers.console.files.FileService.upload_file") - def test_should_handle_file_too_large_error(self, mock_upload): - """Test that service FileTooLargeError is properly converted""" - mock_upload.side_effect = ServiceFileTooLargeError("File too large") - - try: - mock_upload(filename="test.txt", content=b"data", mimetype="text/plain", user=None, source=None) - except ServiceFileTooLargeError as e: - # Simulate the error conversion in FileApi.post() - with pytest.raises(FileTooLargeError): - raise FileTooLargeError(e.description) - - @patch("controllers.console.files.FileService.upload_file") - def test_should_handle_unsupported_file_type_error(self, mock_upload): - """Test that service UnsupportedFileTypeError is properly converted""" - mock_upload.side_effect = ServiceUnsupportedFileTypeError() - - try: - mock_upload( - filename="test.exe", content=b"data", mimetype="application/octet-stream", user=None, source=None - ) - except ServiceUnsupportedFileTypeError: - # Simulate the error conversion in FileApi.post() - with pytest.raises(UnsupportedFileTypeError): - raise UnsupportedFileTypeError() - - # Test 5: File type security + # Test 4: File type security def test_should_identify_dangerous_file_extensions(self): """Test detection of potentially dangerous file extensions""" dangerous_extensions = [ @@ -201,7 +169,7 @@ class TestFileUploadSecurity: parts = filename.split(".") assert len(parts) > 2, f"Filename {filename} should have multiple extensions" - # Test 6: Configuration validation + # Test 5: Configuration validation def test_upload_configuration_structure(self): """Test that upload configuration has correct structure""" # Simulate the configuration returned by FileApi.get()