refactor: port api/fields/file_fields.py (#30638)

This commit is contained in:
Asuka Minato 2026-01-06 23:55:58 +09:00 committed by GitHub
parent 55de731f9c
commit 0294555893
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 257 additions and 162 deletions

View File

@ -1,7 +1,7 @@
from typing import Literal
from flask import request
from flask_restx import Resource, marshal_with
from flask_restx import Resource
from werkzeug.exceptions import Forbidden
import services
@ -15,18 +15,21 @@ from controllers.common.errors import (
TooManyFilesError,
UnsupportedFileTypeError,
)
from controllers.common.schema import register_schema_models
from controllers.console.wraps import (
account_initialization_required,
cloud_edition_billing_resource_check,
setup_required,
)
from extensions.ext_database import db
from fields.file_fields import file_fields, upload_config_fields
from fields.file_fields import FileResponse, UploadConfig
from libs.login import current_account_with_tenant, login_required
from services.file_service import FileService
from . import console_ns
register_schema_models(console_ns, UploadConfig, FileResponse)
PREVIEW_WORDS_LIMIT = 3000
@ -35,26 +38,27 @@ class FileApi(Resource):
@setup_required
@login_required
@account_initialization_required
@marshal_with(upload_config_fields)
@console_ns.response(200, "Success", console_ns.models[UploadConfig.__name__])
def get(self):
return {
"file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,
"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,
"video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
"audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
"workflow_file_upload_limit": dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
"image_file_batch_limit": dify_config.IMAGE_FILE_BATCH_LIMIT,
"single_chunk_attachment_limit": dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT,
"attachment_image_file_size_limit": dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT,
}, 200
config = UploadConfig(
file_size_limit=dify_config.UPLOAD_FILE_SIZE_LIMIT,
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,
video_file_size_limit=dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
audio_file_size_limit=dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
workflow_file_upload_limit=dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
image_file_batch_limit=dify_config.IMAGE_FILE_BATCH_LIMIT,
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
@setup_required
@login_required
@account_initialization_required
@marshal_with(file_fields)
@cloud_edition_billing_resource_check("documents")
@console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__])
def post(self):
current_user, _ = current_account_with_tenant()
source_str = request.form.get("source")
@ -90,7 +94,8 @@ class FileApi(Resource):
except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
raise BlockedFileExtensionError(blocked_extension_error.description)
return upload_file, 201
response = FileResponse.model_validate(upload_file, from_attributes=True)
return response.model_dump(mode="json"), 201
@console_ns.route("/files/<uuid:file_id>/preview")

View File

@ -1,7 +1,7 @@
import urllib.parse
import httpx
from flask_restx import Resource, marshal_with
from flask_restx import Resource
from pydantic import BaseModel, Field
import services
@ -11,19 +11,22 @@ from controllers.common.errors import (
RemoteFileUploadError,
UnsupportedFileTypeError,
)
from controllers.common.schema import register_schema_models
from core.file import helpers as file_helpers
from core.helper import ssrf_proxy
from extensions.ext_database import db
from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
from libs.login import current_account_with_tenant
from services.file_service import FileService
from . import console_ns
register_schema_models(console_ns, RemoteFileInfo, FileWithSignedUrl)
@console_ns.route("/remote-files/<path:url>")
class RemoteFileInfoApi(Resource):
@marshal_with(remote_file_info_fields)
@console_ns.response(200, "Remote file info", console_ns.models[RemoteFileInfo.__name__])
def get(self, url):
decoded_url = urllib.parse.unquote(url)
resp = ssrf_proxy.head(decoded_url)
@ -31,10 +34,11 @@ class RemoteFileInfoApi(Resource):
# failed back to get method
resp = ssrf_proxy.get(decoded_url, timeout=3)
resp.raise_for_status()
return {
"file_type": resp.headers.get("Content-Type", "application/octet-stream"),
"file_length": int(resp.headers.get("Content-Length", 0)),
}
info = RemoteFileInfo(
file_type=resp.headers.get("Content-Type", "application/octet-stream"),
file_length=int(resp.headers.get("Content-Length", 0)),
)
return info.model_dump(mode="json")
class RemoteFileUploadPayload(BaseModel):
@ -50,7 +54,7 @@ console_ns.schema_model(
@console_ns.route("/remote-files/upload")
class RemoteFileUploadApi(Resource):
@console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__])
@marshal_with(file_fields_with_signed_url)
@console_ns.response(201, "Remote file uploaded", console_ns.models[FileWithSignedUrl.__name__])
def post(self):
args = RemoteFileUploadPayload.model_validate(console_ns.payload)
url = args.url
@ -85,13 +89,14 @@ class RemoteFileUploadApi(Resource):
except services.errors.file.UnsupportedFileTypeError:
raise UnsupportedFileTypeError()
return {
"id": upload_file.id,
"name": upload_file.name,
"size": upload_file.size,
"extension": upload_file.extension,
"url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
"mime_type": upload_file.mime_type,
"created_by": upload_file.created_by,
"created_at": upload_file.created_at,
}, 201
payload = FileWithSignedUrl(
id=upload_file.id,
name=upload_file.name,
size=upload_file.size,
extension=upload_file.extension,
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
mime_type=upload_file.mime_type,
created_by=upload_file.created_by,
created_at=int(upload_file.created_at.timestamp()),
)
return payload.model_dump(mode="json"), 201

View File

@ -4,18 +4,18 @@ from flask import request
from flask_restx import Resource
from flask_restx.api import HTTPStatus
from pydantic import BaseModel, Field
from werkzeug.datastructures import FileStorage
from werkzeug.exceptions import Forbidden
import services
from core.file.helpers import verify_plugin_file_signature
from core.tools.tool_file_manager import ToolFileManager
from fields.file_fields import build_file_model
from fields.file_fields import FileResponse
from ..common.errors import (
FileTooLargeError,
UnsupportedFileTypeError,
)
from ..common.schema import register_schema_models
from ..console.wraps import setup_required
from ..files import files_ns
from ..inner_api.plugin.wraps import get_user
@ -35,6 +35,8 @@ files_ns.schema_model(
PluginUploadQuery.__name__, PluginUploadQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)
)
register_schema_models(files_ns, FileResponse)
@files_ns.route("/upload/for-plugin")
class PluginUploadFileApi(Resource):
@ -51,7 +53,7 @@ class PluginUploadFileApi(Resource):
415: "Unsupported file type",
}
)
@files_ns.marshal_with(build_file_model(files_ns), code=HTTPStatus.CREATED)
@files_ns.response(HTTPStatus.CREATED, "File uploaded", files_ns.models[FileResponse.__name__])
def post(self):
"""Upload a file for plugin usage.
@ -69,7 +71,7 @@ class PluginUploadFileApi(Resource):
"""
args = PluginUploadQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
file: FileStorage | None = request.files.get("file")
file = request.files.get("file")
if file is None:
raise Forbidden("File is required.")
@ -80,8 +82,8 @@ class PluginUploadFileApi(Resource):
user_id = args.user_id
user = get_user(tenant_id, user_id)
filename: str | None = file.filename
mimetype: str | None = file.mimetype
filename = file.filename
mimetype = file.mimetype
if not filename or not mimetype:
raise Forbidden("Invalid request.")
@ -111,22 +113,22 @@ class PluginUploadFileApi(Resource):
preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension)
# Create a dictionary with all the necessary attributes
result = {
"id": tool_file.id,
"user_id": tool_file.user_id,
"tenant_id": tool_file.tenant_id,
"conversation_id": tool_file.conversation_id,
"file_key": tool_file.file_key,
"mimetype": tool_file.mimetype,
"original_url": tool_file.original_url,
"name": tool_file.name,
"size": tool_file.size,
"mime_type": mimetype,
"extension": extension,
"preview_url": preview_url,
}
result = FileResponse(
id=tool_file.id,
name=tool_file.name,
size=tool_file.size,
extension=extension,
mime_type=mimetype,
preview_url=preview_url,
source_url=tool_file.original_url,
original_url=tool_file.original_url,
user_id=tool_file.user_id,
tenant_id=tool_file.tenant_id,
conversation_id=tool_file.conversation_id,
file_key=tool_file.file_key,
)
return result, 201
return result.model_dump(mode="json"), 201
except services.errors.file.FileTooLargeError as file_too_large_error:
raise FileTooLargeError(file_too_large_error.description)
except services.errors.file.UnsupportedFileTypeError:

View File

@ -10,13 +10,16 @@ from controllers.common.errors import (
TooManyFilesError,
UnsupportedFileTypeError,
)
from controllers.common.schema import register_schema_models
from controllers.service_api import service_api_ns
from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
from extensions.ext_database import db
from fields.file_fields import build_file_model
from fields.file_fields import FileResponse
from models import App, EndUser
from services.file_service import FileService
register_schema_models(service_api_ns, FileResponse)
@service_api_ns.route("/files/upload")
class FileApi(Resource):
@ -31,8 +34,8 @@ class FileApi(Resource):
415: "Unsupported file type",
}
)
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM))
@service_api_ns.marshal_with(build_file_model(service_api_ns), code=HTTPStatus.CREATED)
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM)) # type: ignore
@service_api_ns.response(HTTPStatus.CREATED, "File uploaded", service_api_ns.models[FileResponse.__name__])
def post(self, app_model: App, end_user: EndUser):
"""Upload a file for use in conversations.
@ -64,4 +67,5 @@ class FileApi(Resource):
except services.errors.file.UnsupportedFileTypeError:
raise UnsupportedFileTypeError()
return upload_file, 201
response = FileResponse.model_validate(upload_file, from_attributes=True)
return response.model_dump(mode="json"), 201

View File

@ -1,5 +1,4 @@
from flask import request
from flask_restx import marshal_with
import services
from controllers.common.errors import (
@ -9,12 +8,15 @@ from controllers.common.errors import (
TooManyFilesError,
UnsupportedFileTypeError,
)
from controllers.common.schema import register_schema_models
from controllers.web import web_ns
from controllers.web.wraps import WebApiResource
from extensions.ext_database import db
from fields.file_fields import build_file_model
from fields.file_fields import FileResponse
from services.file_service import FileService
register_schema_models(web_ns, FileResponse)
@web_ns.route("/files/upload")
class FileApi(WebApiResource):
@ -28,7 +30,7 @@ class FileApi(WebApiResource):
415: "Unsupported file type",
}
)
@marshal_with(build_file_model(web_ns))
@web_ns.response(201, "File uploaded successfully", web_ns.models[FileResponse.__name__])
def post(self, app_model, end_user):
"""Upload a file for use in web applications.
@ -81,4 +83,5 @@ class FileApi(WebApiResource):
except services.errors.file.UnsupportedFileTypeError:
raise UnsupportedFileTypeError()
return upload_file, 201
response = FileResponse.model_validate(upload_file, from_attributes=True)
return response.model_dump(mode="json"), 201

View File

@ -1,7 +1,6 @@
import urllib.parse
import httpx
from flask_restx import marshal_with
from pydantic import BaseModel, Field, HttpUrl
import services
@ -14,7 +13,7 @@ from controllers.common.errors import (
from core.file import helpers as file_helpers
from core.helper import ssrf_proxy
from extensions.ext_database import db
from fields.file_fields import build_file_with_signed_url_model, build_remote_file_info_model
from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
from services.file_service import FileService
from ..common.schema import register_schema_models
@ -26,7 +25,7 @@ class RemoteFileUploadPayload(BaseModel):
url: HttpUrl = Field(description="Remote file URL")
register_schema_models(web_ns, RemoteFileUploadPayload)
register_schema_models(web_ns, RemoteFileUploadPayload, RemoteFileInfo, FileWithSignedUrl)
@web_ns.route("/remote-files/<path:url>")
@ -41,7 +40,7 @@ class RemoteFileInfoApi(WebApiResource):
500: "Failed to fetch remote file",
}
)
@marshal_with(build_remote_file_info_model(web_ns))
@web_ns.response(200, "Remote file info", web_ns.models[RemoteFileInfo.__name__])
def get(self, app_model, end_user, url):
"""Get information about a remote file.
@ -65,10 +64,11 @@ class RemoteFileInfoApi(WebApiResource):
# failed back to get method
resp = ssrf_proxy.get(decoded_url, timeout=3)
resp.raise_for_status()
return {
"file_type": resp.headers.get("Content-Type", "application/octet-stream"),
"file_length": int(resp.headers.get("Content-Length", -1)),
}
info = RemoteFileInfo(
file_type=resp.headers.get("Content-Type", "application/octet-stream"),
file_length=int(resp.headers.get("Content-Length", -1)),
)
return info.model_dump(mode="json")
@web_ns.route("/remote-files/upload")
@ -84,7 +84,7 @@ class RemoteFileUploadApi(WebApiResource):
500: "Failed to fetch remote file",
}
)
@marshal_with(build_file_with_signed_url_model(web_ns))
@web_ns.response(201, "Remote file uploaded", web_ns.models[FileWithSignedUrl.__name__])
def post(self, app_model, end_user):
"""Upload a file from a remote URL.
@ -139,13 +139,14 @@ class RemoteFileUploadApi(WebApiResource):
except services.errors.file.UnsupportedFileTypeError:
raise UnsupportedFileTypeError
return {
"id": upload_file.id,
"name": upload_file.name,
"size": upload_file.size,
"extension": upload_file.extension,
"url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
"mime_type": upload_file.mime_type,
"created_by": upload_file.created_by,
"created_at": upload_file.created_at,
}, 201
payload1 = FileWithSignedUrl(
id=upload_file.id,
name=upload_file.name,
size=upload_file.size,
extension=upload_file.extension,
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
mime_type=upload_file.mime_type,
created_by=upload_file.created_by,
created_at=int(upload_file.created_at.timestamp()),
)
return payload1.model_dump(mode="json"), 201

View File

@ -1,93 +1,85 @@
from flask_restx import Namespace, fields
from __future__ import annotations
from libs.helper import TimestampField
from datetime import datetime
upload_config_fields = {
"file_size_limit": fields.Integer,
"batch_count_limit": fields.Integer,
"image_file_size_limit": fields.Integer,
"video_file_size_limit": fields.Integer,
"audio_file_size_limit": fields.Integer,
"workflow_file_upload_limit": fields.Integer,
"image_file_batch_limit": fields.Integer,
"single_chunk_attachment_limit": fields.Integer,
}
from pydantic import BaseModel, ConfigDict, field_validator
def build_upload_config_model(api_or_ns: Namespace):
"""Build the upload config model for the API or Namespace.
Args:
api_or_ns: Flask-RestX Api or Namespace instance
Returns:
The registered model
"""
return api_or_ns.model("UploadConfig", upload_config_fields)
class ResponseModel(BaseModel):
model_config = ConfigDict(
from_attributes=True,
extra="ignore",
populate_by_name=True,
serialize_by_alias=True,
protected_namespaces=(),
)
file_fields = {
"id": fields.String,
"name": fields.String,
"size": fields.Integer,
"extension": fields.String,
"mime_type": fields.String,
"created_by": fields.String,
"created_at": TimestampField,
"preview_url": fields.String,
"source_url": fields.String,
}
def _to_timestamp(value: datetime | int | None) -> int | None:
if isinstance(value, datetime):
return int(value.timestamp())
return value
def build_file_model(api_or_ns: Namespace):
"""Build the file model for the API or Namespace.
Args:
api_or_ns: Flask-RestX Api or Namespace instance
Returns:
The registered model
"""
return api_or_ns.model("File", file_fields)
class UploadConfig(ResponseModel):
file_size_limit: int
batch_count_limit: int
file_upload_limit: int | None = None
image_file_size_limit: int
video_file_size_limit: int
audio_file_size_limit: int
workflow_file_upload_limit: int
image_file_batch_limit: int
single_chunk_attachment_limit: int
attachment_image_file_size_limit: int | None = None
remote_file_info_fields = {
"file_type": fields.String(attribute="file_type"),
"file_length": fields.Integer(attribute="file_length"),
}
class FileResponse(ResponseModel):
id: str
name: str
size: int
extension: str | None = None
mime_type: str | None = None
created_by: str | None = None
created_at: int | None = None
preview_url: str | None = None
source_url: str | None = None
original_url: str | None = None
user_id: str | None = None
tenant_id: str | None = None
conversation_id: str | None = None
file_key: str | None = None
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
def build_remote_file_info_model(api_or_ns: Namespace):
"""Build the remote file info model for the API or Namespace.
Args:
api_or_ns: Flask-RestX Api or Namespace instance
Returns:
The registered model
"""
return api_or_ns.model("RemoteFileInfo", remote_file_info_fields)
class RemoteFileInfo(ResponseModel):
file_type: str
file_length: int
file_fields_with_signed_url = {
"id": fields.String,
"name": fields.String,
"size": fields.Integer,
"extension": fields.String,
"url": fields.String,
"mime_type": fields.String,
"created_by": fields.String,
"created_at": TimestampField,
}
class FileWithSignedUrl(ResponseModel):
id: str
name: str
size: int
extension: str | None = None
url: str | None = None
mime_type: str | None = None
created_by: str | None = None
created_at: int | None = None
@field_validator("created_at", mode="before")
@classmethod
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
return _to_timestamp(value)
def build_file_with_signed_url_model(api_or_ns: Namespace):
"""Build the file with signed URL model for the API or Namespace.
Args:
api_or_ns: Flask-RestX Api or Namespace instance
Returns:
The registered model
"""
return api_or_ns.model("FileWithSignedUrl", file_fields_with_signed_url)
__all__ = [
"FileResponse",
"FileWithSignedUrl",
"RemoteFileInfo",
"UploadConfig",
]

View File

@ -1,7 +1,9 @@
import builtins
import io
from unittest.mock import patch
import pytest
from flask.views import MethodView
from werkzeug.exceptions import Forbidden
from controllers.common.errors import (
@ -14,6 +16,9 @@ from controllers.common.errors import (
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]
class TestFileUploadSecurity:
"""Test file upload security logic without complex framework setup"""
@ -128,7 +133,7 @@ class TestFileUploadSecurity:
# Test passes if no exception is raised
# Test 4: Service error handling
@patch("services.file_service.FileService.upload_file")
@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")
@ -140,7 +145,7 @@ class TestFileUploadSecurity:
with pytest.raises(FileTooLargeError):
raise FileTooLargeError(e.description)
@patch("services.file_service.FileService.upload_file")
@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()

View File

@ -0,0 +1,78 @@
from __future__ import annotations
from datetime import datetime
from types import SimpleNamespace
from fields.file_fields import FileResponse, FileWithSignedUrl, RemoteFileInfo, UploadConfig
def test_file_response_serializes_datetime() -> None:
created_at = datetime(2024, 1, 1, 12, 0, 0)
file_obj = SimpleNamespace(
id="file-1",
name="example.txt",
size=1024,
extension="txt",
mime_type="text/plain",
created_by="user-1",
created_at=created_at,
preview_url="https://preview",
source_url="https://source",
original_url="https://origin",
user_id="user-1",
tenant_id="tenant-1",
conversation_id="conv-1",
file_key="key-1",
)
serialized = FileResponse.model_validate(file_obj, from_attributes=True).model_dump(mode="json")
assert serialized["id"] == "file-1"
assert serialized["created_at"] == int(created_at.timestamp())
assert serialized["preview_url"] == "https://preview"
assert serialized["source_url"] == "https://source"
assert serialized["original_url"] == "https://origin"
assert serialized["user_id"] == "user-1"
assert serialized["tenant_id"] == "tenant-1"
assert serialized["conversation_id"] == "conv-1"
assert serialized["file_key"] == "key-1"
def test_file_with_signed_url_builds_payload() -> None:
payload = FileWithSignedUrl(
id="file-2",
name="remote.pdf",
size=2048,
extension="pdf",
url="https://signed",
mime_type="application/pdf",
created_by="user-2",
created_at=datetime(2024, 1, 2, 0, 0, 0),
)
dumped = payload.model_dump(mode="json")
assert dumped["url"] == "https://signed"
assert dumped["created_at"] == int(datetime(2024, 1, 2, 0, 0, 0).timestamp())
def test_remote_file_info_and_upload_config() -> None:
info = RemoteFileInfo(file_type="text/plain", file_length=123)
assert info.model_dump(mode="json") == {"file_type": "text/plain", "file_length": 123}
config = UploadConfig(
file_size_limit=1,
batch_count_limit=2,
file_upload_limit=3,
image_file_size_limit=4,
video_file_size_limit=5,
audio_file_size_limit=6,
workflow_file_upload_limit=7,
image_file_batch_limit=8,
single_chunk_attachment_limit=9,
attachment_image_file_size_limit=10,
)
dumped = config.model_dump(mode="json")
assert dumped["file_upload_limit"] == 3
assert dumped["attachment_image_file_size_limit"] == 10