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

@@ -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/<path:url>")
-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/<path:url>",
+    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()),
+    )

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

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()),
+    }