refactor: Migrate part of the console basic API module to Flask-RESTX (#24732)

Signed-off-by: -LAN- <laipz8200@outlook.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com> Co-authored-by: -LAN- <laipz8200@outlook.com>
2025-09-10 12:15:47 +08:00
parent 26a9abef64
commit b51c724a94
27 changed files with 917 additions and 206 deletions
--- a/api/controllers/console/auth/activate.py
+++ b/api/controllers/console/auth/activate.py
@@ -1,8 +1,8 @@
 from flask import request
-from flask_restx import Resource, reqparse
+from flask_restx import Resource, fields, reqparse

 from constants.languages import supported_language
-from controllers.console import api
+from controllers.console import api, console_ns
 from controllers.console.error import AlreadyActivateError
 from extensions.ext_database import db
 from libs.datetime_utils import naive_utc_now
@@ -10,14 +10,36 @@ from libs.helper import StrLen, email, extract_remote_ip, timezone
 from models.account import AccountStatus
 from services.account_service import AccountService, RegisterService

+active_check_parser = reqparse.RequestParser()
+active_check_parser.add_argument(
+    "workspace_id", type=str, required=False, nullable=True, location="args", help="Workspace ID"
+)
+active_check_parser.add_argument(
+    "email", type=email, required=False, nullable=True, location="args", help="Email address"
+)
+active_check_parser.add_argument(
+    "token", type=str, required=True, nullable=False, location="args", help="Activation token"
+)

+
+@console_ns.route("/activate/check")
 class ActivateCheckApi(Resource):
+    @api.doc("check_activation_token")
+    @api.doc(description="Check if activation token is valid")
+    @api.expect(active_check_parser)
+    @api.response(
+        200,
+        "Success",
+        api.model(
+            "ActivationCheckResponse",
+            {
+                "is_valid": fields.Boolean(description="Whether token is valid"),
+                "data": fields.Raw(description="Activation data if valid"),
+            },
+        ),
+    )
    def get(self):
-        parser = reqparse.RequestParser()
-        parser.add_argument("workspace_id", type=str, required=False, nullable=True, location="args")
-        parser.add_argument("email", type=email, required=False, nullable=True, location="args")
-        parser.add_argument("token", type=str, required=True, nullable=False, location="args")
-        args = parser.parse_args()
+        args = active_check_parser.parse_args()

        workspaceId = args["workspace_id"]
        reg_email = args["email"]
@@ -38,18 +60,36 @@ class ActivateCheckApi(Resource):
            return {"is_valid": False}


+active_parser = reqparse.RequestParser()
+active_parser.add_argument("workspace_id", type=str, required=False, nullable=True, location="json")
+active_parser.add_argument("email", type=email, required=False, nullable=True, location="json")
+active_parser.add_argument("token", type=str, required=True, nullable=False, location="json")
+active_parser.add_argument("name", type=StrLen(30), required=True, nullable=False, location="json")
+active_parser.add_argument(
+    "interface_language", type=supported_language, required=True, nullable=False, location="json"
+)
+active_parser.add_argument("timezone", type=timezone, required=True, nullable=False, location="json")
+
+
+@console_ns.route("/activate")
 class ActivateApi(Resource):
+    @api.doc("activate_account")
+    @api.doc(description="Activate account with invitation token")
+    @api.expect(active_parser)
+    @api.response(
+        200,
+        "Account activated successfully",
+        api.model(
+            "ActivationResponse",
+            {
+                "result": fields.String(description="Operation result"),
+                "data": fields.Raw(description="Login token data"),
+            },
+        ),
+    )
+    @api.response(400, "Already activated or invalid token")
    def post(self):
-        parser = reqparse.RequestParser()
-        parser.add_argument("workspace_id", type=str, required=False, nullable=True, location="json")
-        parser.add_argument("email", type=email, required=False, nullable=True, location="json")
-        parser.add_argument("token", type=str, required=True, nullable=False, location="json")
-        parser.add_argument("name", type=StrLen(30), required=True, nullable=False, location="json")
-        parser.add_argument(
-            "interface_language", type=supported_language, required=True, nullable=False, location="json"
-        )
-        parser.add_argument("timezone", type=timezone, required=True, nullable=False, location="json")
-        args = parser.parse_args()
+        args = active_parser.parse_args()

        invitation = RegisterService.get_invitation_if_token_valid(args["workspace_id"], args["email"], args["token"])
        if invitation is None:
@@ -70,7 +110,3 @@ class ActivateApi(Resource):
        token_pair = AccountService.login(account, ip_address=extract_remote_ip(request))

        return {"result": "success", "data": token_pair.model_dump()}
-
-
-api.add_resource(ActivateCheckApi, "/activate/check")
-api.add_resource(ActivateApi, "/activate")
--- a/api/controllers/console/auth/data_source_oauth.py
+++ b/api/controllers/console/auth/data_source_oauth.py
@@ -3,11 +3,11 @@ import logging
 import requests
 from flask import current_app, redirect, request
 from flask_login import current_user
-from flask_restx import Resource
+from flask_restx import Resource, fields
 from werkzeug.exceptions import Forbidden

 from configs import dify_config
-from controllers.console import api
+from controllers.console import api, console_ns
 from libs.login import login_required
 from libs.oauth_data_source import NotionOAuth

@@ -28,7 +28,21 @@ def get_oauth_providers():
        return OAUTH_PROVIDERS


+@console_ns.route("/oauth/data-source/<string:provider>")
 class OAuthDataSource(Resource):
+    @api.doc("oauth_data_source")
+    @api.doc(description="Get OAuth authorization URL for data source provider")
+    @api.doc(params={"provider": "Data source provider name (notion)"})
+    @api.response(
+        200,
+        "Authorization URL or internal setup success",
+        api.model(
+            "OAuthDataSourceResponse",
+            {"data": fields.Raw(description="Authorization URL or 'internal' for internal setup")},
+        ),
+    )
+    @api.response(400, "Invalid provider")
+    @api.response(403, "Admin privileges required")
    def get(self, provider: str):
        # The role of the current user in the table must be admin or owner
        if not current_user.is_admin_or_owner:
@@ -49,7 +63,19 @@ class OAuthDataSource(Resource):
            return {"data": auth_url}, 200


+@console_ns.route("/oauth/data-source/callback/<string:provider>")
 class OAuthDataSourceCallback(Resource):
+    @api.doc("oauth_data_source_callback")
+    @api.doc(description="Handle OAuth callback from data source provider")
+    @api.doc(
+        params={
+            "provider": "Data source provider name (notion)",
+            "code": "Authorization code from OAuth provider",
+            "error": "Error message from OAuth provider",
+        }
+    )
+    @api.response(302, "Redirect to console with result")
+    @api.response(400, "Invalid provider")
    def get(self, provider: str):
        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
        with current_app.app_context():
@@ -68,7 +94,19 @@ class OAuthDataSourceCallback(Resource):
            return redirect(f"{dify_config.CONSOLE_WEB_URL}?type=notion&error=Access denied")


+@console_ns.route("/oauth/data-source/binding/<string:provider>")
 class OAuthDataSourceBinding(Resource):
+    @api.doc("oauth_data_source_binding")
+    @api.doc(description="Bind OAuth data source with authorization code")
+    @api.doc(
+        params={"provider": "Data source provider name (notion)", "code": "Authorization code from OAuth provider"}
+    )
+    @api.response(
+        200,
+        "Data source binding success",
+        api.model("OAuthDataSourceBindingResponse", {"result": fields.String(description="Operation result")}),
+    )
+    @api.response(400, "Invalid provider or code")
    def get(self, provider: str):
        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
        with current_app.app_context():
@@ -90,7 +128,17 @@ class OAuthDataSourceBinding(Resource):
            return {"result": "success"}, 200


+@console_ns.route("/oauth/data-source/<string:provider>/<uuid:binding_id>/sync")
 class OAuthDataSourceSync(Resource):
+    @api.doc("oauth_data_source_sync")
+    @api.doc(description="Sync data from OAuth data source")
+    @api.doc(params={"provider": "Data source provider name (notion)", "binding_id": "Data source binding ID"})
+    @api.response(
+        200,
+        "Data source sync success",
+        api.model("OAuthDataSourceSyncResponse", {"result": fields.String(description="Operation result")}),
+    )
+    @api.response(400, "Invalid provider or sync failed")
    @setup_required
    @login_required
    @account_initialization_required
@@ -111,9 +159,3 @@ class OAuthDataSourceSync(Resource):
            return {"error": "OAuth data source process failed"}, 400

        return {"result": "success"}, 200
-
-
-api.add_resource(OAuthDataSource, "/oauth/data-source/<string:provider>")
-api.add_resource(OAuthDataSourceCallback, "/oauth/data-source/callback/<string:provider>")
-api.add_resource(OAuthDataSourceBinding, "/oauth/data-source/binding/<string:provider>")
-api.add_resource(OAuthDataSourceSync, "/oauth/data-source/<string:provider>/<uuid:binding_id>/sync")
--- a/api/controllers/console/auth/forgot_password.py
+++ b/api/controllers/console/auth/forgot_password.py
@@ -2,12 +2,12 @@ import base64
 import secrets

 from flask import request
-from flask_restx import Resource, reqparse
+from flask_restx import Resource, fields, reqparse
 from sqlalchemy import select
 from sqlalchemy.orm import Session

 from constants.languages import languages
-from controllers.console import api
+from controllers.console import api, console_ns
 from controllers.console.auth.error import (
    EmailCodeError,
    EmailPasswordResetLimitError,
@@ -28,7 +28,32 @@ from services.errors.workspace import WorkSpaceNotAllowedCreateError, Workspaces
 from services.feature_service import FeatureService


+@console_ns.route("/forgot-password")
 class ForgotPasswordSendEmailApi(Resource):
+    @api.doc("send_forgot_password_email")
+    @api.doc(description="Send password reset email")
+    @api.expect(
+        api.model(
+            "ForgotPasswordEmailRequest",
+            {
+                "email": fields.String(required=True, description="Email address"),
+                "language": fields.String(description="Language for email (zh-Hans/en-US)"),
+            },
+        )
+    )
+    @api.response(
+        200,
+        "Email sent successfully",
+        api.model(
+            "ForgotPasswordEmailResponse",
+            {
+                "result": fields.String(description="Operation result"),
+                "data": fields.String(description="Reset token"),
+                "code": fields.String(description="Error code if account not found"),
+            },
+        ),
+    )
+    @api.response(400, "Invalid email or rate limit exceeded")
    @setup_required
    @email_password_login_enabled
    def post(self):
@@ -61,7 +86,33 @@ class ForgotPasswordSendEmailApi(Resource):
        return {"result": "success", "data": token}


+@console_ns.route("/forgot-password/validity")
 class ForgotPasswordCheckApi(Resource):
+    @api.doc("check_forgot_password_code")
+    @api.doc(description="Verify password reset code")
+    @api.expect(
+        api.model(
+            "ForgotPasswordCheckRequest",
+            {
+                "email": fields.String(required=True, description="Email address"),
+                "code": fields.String(required=True, description="Verification code"),
+                "token": fields.String(required=True, description="Reset token"),
+            },
+        )
+    )
+    @api.response(
+        200,
+        "Code verified successfully",
+        api.model(
+            "ForgotPasswordCheckResponse",
+            {
+                "is_valid": fields.Boolean(description="Whether code is valid"),
+                "email": fields.String(description="Email address"),
+                "token": fields.String(description="New reset token"),
+            },
+        ),
+    )
+    @api.response(400, "Invalid code or token")
    @setup_required
    @email_password_login_enabled
    def post(self):
@@ -100,7 +151,26 @@ class ForgotPasswordCheckApi(Resource):
        return {"is_valid": True, "email": token_data.get("email"), "token": new_token}


+@console_ns.route("/forgot-password/resets")
 class ForgotPasswordResetApi(Resource):
+    @api.doc("reset_password")
+    @api.doc(description="Reset password with verification token")
+    @api.expect(
+        api.model(
+            "ForgotPasswordResetRequest",
+            {
+                "token": fields.String(required=True, description="Verification token"),
+                "new_password": fields.String(required=True, description="New password"),
+                "password_confirm": fields.String(required=True, description="Password confirmation"),
+            },
+        )
+    )
+    @api.response(
+        200,
+        "Password reset successfully",
+        api.model("ForgotPasswordResetResponse", {"result": fields.String(description="Operation result")}),
+    )
+    @api.response(400, "Invalid token or password mismatch")
    @setup_required
    @email_password_login_enabled
    def post(self):
@@ -172,8 +242,3 @@ class ForgotPasswordResetApi(Resource):
            pass
        except AccountRegisterError:
            raise AccountInFreezeError()
-
-
-api.add_resource(ForgotPasswordSendEmailApi, "/forgot-password")
-api.add_resource(ForgotPasswordCheckApi, "/forgot-password/validity")
-api.add_resource(ForgotPasswordResetApi, "/forgot-password/resets")
--- a/api/controllers/console/auth/oauth.py
+++ b/api/controllers/console/auth/oauth.py
@@ -22,7 +22,7 @@ from services.errors.account import AccountNotFoundError, AccountRegisterError
 from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkSpaceNotFoundError
 from services.feature_service import FeatureService

-from .. import api
+from .. import api, console_ns

 logger = logging.getLogger(__name__)

@@ -50,7 +50,13 @@ def get_oauth_providers():
        return OAUTH_PROVIDERS


+@console_ns.route("/oauth/login/<provider>")
 class OAuthLogin(Resource):
+    @api.doc("oauth_login")
+    @api.doc(description="Initiate OAuth login process")
+    @api.doc(params={"provider": "OAuth provider name (github/google)", "invite_token": "Optional invitation token"})
+    @api.response(302, "Redirect to OAuth authorization URL")
+    @api.response(400, "Invalid provider")
    def get(self, provider: str):
        invite_token = request.args.get("invite_token") or None
        OAUTH_PROVIDERS = get_oauth_providers()
@@ -63,7 +69,19 @@ class OAuthLogin(Resource):
        return redirect(auth_url)


+@console_ns.route("/oauth/authorize/<provider>")
 class OAuthCallback(Resource):
+    @api.doc("oauth_callback")
+    @api.doc(description="Handle OAuth callback and complete login process")
+    @api.doc(
+        params={
+            "provider": "OAuth provider name (github/google)",
+            "code": "Authorization code from OAuth provider",
+            "state": "Optional state parameter (used for invite token)",
+        }
+    )
+    @api.response(302, "Redirect to console with access token")
+    @api.response(400, "OAuth process failed")
    def get(self, provider: str):
        OAUTH_PROVIDERS = get_oauth_providers()
        with current_app.app_context():
@@ -184,7 +202,3 @@ def _generate_account(provider: str, user_info: OAuthUserInfo):
    AccountService.link_account_integrate(provider, user_info.id, account)

    return account
-
-
-api.add_resource(OAuthLogin, "/oauth/login/<provider>")
-api.add_resource(OAuthCallback, "/oauth/authorize/<provider>")