From ba568a634d107fe3329393627183e1346ac136fd Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Jan 2026 19:32:20 +0900 Subject: [PATCH] refactor: api/controllers/console/remote_files.py to ov3 (#31466) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/controllers/console/remote_files.py | 127 ++++++++---------- api/extensions/ext_fastopenapi.py | 3 +- .../console/test_fastopenapi_remote_files.py | 92 +++++++++++++ 3 files changed, 152 insertions(+), 70 deletions(-) create mode 100644 api/tests/unit_tests/controllers/console/test_fastopenapi_remote_files.py diff --git a/api/controllers/console/remote_files.py b/api/controllers/console/remote_files.py index 70c7b80ffa..88a9ce3a79 100644 --- a/api/controllers/console/remote_files.py +++ b/api/controllers/console/remote_files.py @@ -1,7 +1,6 @@ import urllib.parse import httpx -from flask_restx import Resource from pydantic import BaseModel, Field import services @@ -11,7 +10,7 @@ from controllers.common.errors import ( RemoteFileUploadError, UnsupportedFileTypeError, ) -from controllers.common.schema import register_schema_models +from controllers.fastopenapi import console_router from core.file import helpers as file_helpers from core.helper import ssrf_proxy from extensions.ext_database import db @@ -19,84 +18,74 @@ 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/") -class RemoteFileInfoApi(Resource): - @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) - if resp.status_code != httpx.codes.OK: - # failed back to get method - resp = ssrf_proxy.get(decoded_url, timeout=3) - resp.raise_for_status() - 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): url: str = Field(..., description="URL to fetch") -console_ns.schema_model( - RemoteFileUploadPayload.__name__, - RemoteFileUploadPayload.model_json_schema(ref_template="#/definitions/{model}"), +@console_router.get( + "/remote-files/", + response_model=RemoteFileInfo, + tags=["console"], ) +def get_remote_file_info(url: str) -> RemoteFileInfo: + decoded_url = urllib.parse.unquote(url) + resp = ssrf_proxy.head(decoded_url) + if resp.status_code != httpx.codes.OK: + resp = ssrf_proxy.get(decoded_url, timeout=3) + resp.raise_for_status() + return RemoteFileInfo( + file_type=resp.headers.get("Content-Type", "application/octet-stream"), + file_length=int(resp.headers.get("Content-Length", 0)), + ) -@console_ns.route("/remote-files/upload") -class RemoteFileUploadApi(Resource): - @console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__]) - @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 +@console_router.post( + "/remote-files/upload", + response_model=FileWithSignedUrl, + tags=["console"], + status_code=201, +) +def upload_remote_file(payload: RemoteFileUploadPayload) -> FileWithSignedUrl: + url = payload.url - try: - resp = ssrf_proxy.head(url=url) - if resp.status_code != httpx.codes.OK: - resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True) - if resp.status_code != httpx.codes.OK: - raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}") - except httpx.RequestError as e: - raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}") + try: + resp = ssrf_proxy.head(url=url) + if resp.status_code != httpx.codes.OK: + resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True) + if resp.status_code != httpx.codes.OK: + raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}") + except httpx.RequestError as e: + raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}") - file_info = helpers.guess_file_info_from_response(resp) + file_info = helpers.guess_file_info_from_response(resp) - if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size): - raise FileTooLargeError + if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size): + raise FileTooLargeError - content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content + content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content - try: - user, _ = current_account_with_tenant() - upload_file = FileService(db.engine).upload_file( - filename=file_info.filename, - content=content, - mimetype=file_info.mimetype, - user=user, - source_url=url, - ) - except services.errors.file.FileTooLargeError as file_too_large_error: - raise FileTooLargeError(file_too_large_error.description) - except services.errors.file.UnsupportedFileTypeError: - raise UnsupportedFileTypeError() - - 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()), + try: + user, _ = current_account_with_tenant() + upload_file = FileService(db.engine).upload_file( + filename=file_info.filename, + content=content, + mimetype=file_info.mimetype, + user=user, + source_url=url, ) - return payload.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: + raise UnsupportedFileTypeError() + + return 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()), + ) diff --git a/api/extensions/ext_fastopenapi.py b/api/extensions/ext_fastopenapi.py index e6c1bc6bee..719456803a 100644 --- a/api/extensions/ext_fastopenapi.py +++ b/api/extensions/ext_fastopenapi.py @@ -28,9 +28,10 @@ def init_app(app: DifyApp) -> None: # Ensure route decorators are evaluated. import controllers.console.ping as ping_module - from controllers.console import setup + from controllers.console import remote_files, setup _ = ping_module + _ = remote_files _ = setup router.include_router(console_router, prefix="/console/api") diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_remote_files.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_remote_files.py new file mode 100644 index 0000000000..cb2604cf1c --- /dev/null +++ b/api/tests/unit_tests/controllers/console/test_fastopenapi_remote_files.py @@ -0,0 +1,92 @@ +import builtins +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import patch + +import httpx +import pytest +from flask import Flask +from flask.views import MethodView + +from extensions import ext_fastopenapi + +if not hasattr(builtins, "MethodView"): + builtins.MethodView = MethodView # type: ignore[attr-defined] + + +@pytest.fixture +def app() -> Flask: + app = Flask(__name__) + app.config["TESTING"] = True + return app + + +def test_console_remote_files_fastopenapi_get_info(app: Flask): + ext_fastopenapi.init_app(app) + + response = httpx.Response( + 200, + request=httpx.Request("HEAD", "http://example.com/file.txt"), + headers={"Content-Type": "text/plain", "Content-Length": "10"}, + ) + + with patch("controllers.console.remote_files.ssrf_proxy.head", return_value=response): + client = app.test_client() + encoded_url = "http%3A%2F%2Fexample.com%2Ffile.txt" + resp = client.get(f"/console/api/remote-files/{encoded_url}") + + assert resp.status_code == 200 + assert resp.get_json() == {"file_type": "text/plain", "file_length": 10} + + +def test_console_remote_files_fastopenapi_upload(app: Flask): + ext_fastopenapi.init_app(app) + + head_response = httpx.Response( + 200, + request=httpx.Request("GET", "http://example.com/file.txt"), + content=b"hello", + ) + file_info = SimpleNamespace( + extension="txt", + size=5, + filename="file.txt", + mimetype="text/plain", + ) + uploaded = SimpleNamespace( + id="file-id", + name="file.txt", + size=5, + extension="txt", + mime_type="text/plain", + created_by="user-id", + created_at=datetime(2024, 1, 1), + ) + + with ( + patch("controllers.console.remote_files.db", new=SimpleNamespace(engine=object())), + patch("controllers.console.remote_files.ssrf_proxy.head", return_value=head_response), + patch("controllers.console.remote_files.helpers.guess_file_info_from_response", return_value=file_info), + patch("controllers.console.remote_files.FileService.is_file_size_within_limit", return_value=True), + patch("controllers.console.remote_files.FileService.__init__", return_value=None), + patch("controllers.console.remote_files.current_account_with_tenant", return_value=(object(), "tenant-id")), + patch("controllers.console.remote_files.FileService.upload_file", return_value=uploaded), + patch("controllers.console.remote_files.file_helpers.get_signed_file_url", return_value="signed-url"), + ): + client = app.test_client() + resp = client.post( + "/console/api/remote-files/upload", + json={"url": "http://example.com/file.txt"}, + ) + + assert resp.status_code == 201 + assert resp.get_json() == { + "id": "file-id", + "name": "file.txt", + "size": 5, + "extension": "txt", + "url": "signed-url", + "mime_type": "text/plain", + "created_by": "user-id", + "created_at": int(uploaded.created_at.timestamp()), + }