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

refactor(console): inject file service (#41603)

Browse Source
This commit is contained in:
非法操作 2026-09-03 04:42:30 +00:00 committed by GitHub
parent 5f4b1f867d
commit 69484a07e1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 72 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
api/controllers/console/files.py
View File

@ -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/<uuid:file_id>/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)},
)

54
api/tests/unit_tests/controllers/console/test_files.py
View File

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

36
api/tests/unit_tests/controllers/console/test_files_security.py
View File

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