refactor: better error handler (#24422)

Signed-off-by: -LAN- <laipz8200@outlook.com>

api/extensions/ext_blueprints.py

@@ -29,7 +29,6 @@ def init_app(app: DifyApp):
         methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
         expose_headers=["X-Version", "X-Env"],
     )
-
     app.register_blueprint(web_bp)
 
     CORS(
@@ -40,10 +39,13 @@ def init_app(app: DifyApp):
         methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
         expose_headers=["X-Version", "X-Env"],
     )
-
     app.register_blueprint(console_app_bp)
 
-    CORS(files_bp, allow_headers=["Content-Type"], methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"])
+    CORS(
+        files_bp,
+        allow_headers=["Content-Type"],
+        methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
+    )
     app.register_blueprint(files_bp)
 
     app.register_blueprint(inner_api_bp)

api/extensions/ext_login.py

@@ -20,6 +20,10 @@ login_manager = flask_login.LoginManager()
 @login_manager.request_loader
 def load_user_from_request(request_from_flask_login):
     """Load user based on the request."""
+    # Skip authentication for documentation endpoints
+    if request.path.endswith("/docs") or request.path.endswith("/swagger.json"):
+        return None
+
     auth_header = request.headers.get("Authorization", "")
     auth_token: str | None = None
     if auth_header:

api/libs/external_api.py

@@ -16,98 +16,124 @@ def http_status_message(code):
     return HTTP_STATUS_CODES.get(code, "")
 
 
-class ExternalApi(Api):
+def register_external_error_handlers(api: Api) -> None:
-    def handle_error(self, e):
+    """Register error handlers for the API using decorators.
-        """Error handler for the API transforms a raised exception into a Flask
-        response, with the appropriate HTTP status code and body.
 
-        :param e: the raised Exception object
+    :param api: The Flask-RestX Api instance
-        :type e: Exception
+    """
 
-        """
+    @api.errorhandler(HTTPException)
+    def handle_http_exception(e: HTTPException):
+        """Handle HTTP exceptions."""
         got_request_exception.send(current_app, exception=e)
 
+        if e.response is not None:
+            return e.get_response()
+
         headers = Headers()
-        if isinstance(e, HTTPException):
+        status_code = e.code
-            if e.response is not None:
+        default_data = {
-                resp = e.get_response()
+            "code": re.sub(r"(?<!^)(?=[A-Z])", "_", type(e).__name__).lower(),
-                return resp
+            "message": getattr(e, "description", http_status_message(status_code)),
+            "status": status_code,
+        }
 
-            status_code = e.code
+        if (
-            default_data = {
+            default_data["message"]
-                "code": re.sub(r"(?<!^)(?=[A-Z])", "_", type(e).__name__).lower(),
+            and default_data["message"] == "Failed to decode JSON object: Expecting value: line 1 column 1 (char 0)"
-                "message": getattr(e, "description", http_status_message(status_code)),
+        ):
-                "status": status_code,
+            default_data["message"] = "Invalid JSON payload received or JSON payload is empty."
-            }
 
-            if (
+        headers = e.get_response().headers
-                default_data["message"]
-                and default_data["message"] == "Failed to decode JSON object: Expecting value: line 1 column 1 (char 0)"
-            ):
-                default_data["message"] = "Invalid JSON payload received or JSON payload is empty."
 
-            headers = e.get_response().headers
+        # Handle specific status codes
-        elif isinstance(e, ValueError):
+        if status_code == 406 and api.default_mediatype is None:
-            status_code = 400
+            supported_mediatypes = list(api.representations.keys())
-            default_data = {
+            fallback_mediatype = supported_mediatypes[0] if supported_mediatypes else "text/plain"
-                "code": "invalid_param",
+            data = {"code": "not_acceptable", "message": default_data.get("message")}
-                "message": str(e),
+            resp = api.make_response(data, status_code, headers, fallback_mediatype=fallback_mediatype)
-                "status": status_code,
+        elif status_code == 400:
-            }
+            if isinstance(default_data.get("message"), dict):
-        elif isinstance(e, AppInvokeQuotaExceededError):
+                param_key, param_value = list(default_data.get("message", {}).items())[0]
-            status_code = 429
+                data = {"code": "invalid_param", "message": param_value, "params": param_key}
-            default_data = {
+            else:
-                "code": "too_many_requests",
+                data = default_data
-                "message": str(e),
+                if "code" not in data:
-                "status": status_code,
+                    data["code"] = "unknown"
-            }
+            resp = api.make_response(data, status_code, headers)
         else:
-            status_code = 500
+            data = default_data
-            default_data = {
+            if "code" not in data:
-                "message": http_status_message(status_code),
+                data["code"] = "unknown"
-            }
+            resp = api.make_response(data, status_code, headers)
 
-        # Werkzeug exceptions generate a content-length header which is added
+        if status_code == 401:
-        # to the response in addition to the actual content-length header
+            resp = api.unauthorized(resp)
-        # https://github.com/flask-restful/flask-restful/issues/534
+
+        # Remove duplicate Content-Length header
         remove_headers = ("Content-Length",)
-
         for header in remove_headers:
             headers.pop(header, None)
 
+        return resp
+
+    @api.errorhandler(ValueError)
+    def handle_value_error(e: ValueError):
+        """Handle ValueError exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        status_code = 400
+        data = {
+            "code": "invalid_param",
+            "message": str(e),
+            "status": status_code,
+        }
+        return api.make_response(data, status_code)
+
+    @api.errorhandler(AppInvokeQuotaExceededError)
+    def handle_quota_exceeded(e: AppInvokeQuotaExceededError):
+        """Handle AppInvokeQuotaExceededError exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        status_code = 429
+        data = {
+            "code": "too_many_requests",
+            "message": str(e),
+            "status": status_code,
+        }
+        return api.make_response(data, status_code)
+
+    @api.errorhandler(Exception)
+    def handle_general_exception(e: Exception):
+        """Handle general exceptions."""
+        got_request_exception.send(current_app, exception=e)
+
+        headers = Headers()
+        status_code = 500
+        default_data = {
+            "message": http_status_message(status_code),
+        }
+
         data = getattr(e, "data", default_data)
 
-        # record the exception in the logs when we have a server error of status code: 500
+        # Log server errors
-        if status_code and status_code >= 500:
+        exc_info: Any = sys.exc_info()
-            exc_info: Any = sys.exc_info()
+        if exc_info[1] is None:
-            if exc_info[1] is None:
+            exc_info = None
-                exc_info = None
+        current_app.log_exception(exc_info)
-            current_app.log_exception(exc_info)
 
-        if status_code == 406 and self.default_mediatype is None:
+        if "code" not in data:
-            # if we are handling NotAcceptable (406), make sure that
+            data["code"] = "unknown"
-            # make_response uses a representation we support as the
-            # default mediatype (so that make_response doesn't throw
-            # another NotAcceptable error).
-            supported_mediatypes = list(self.representations.keys())  # only supported application/json
-            fallback_mediatype = supported_mediatypes[0] if supported_mediatypes else "text/plain"
-            data = {"code": "not_acceptable", "message": data.get("message")}
-            resp = self.make_response(data, status_code, headers, fallback_mediatype=fallback_mediatype)
-        elif status_code == 400:
-            if isinstance(data.get("message"), dict):
-                param_key, param_value = list(data.get("message", {}).items())[0]
-                data = {"code": "invalid_param", "message": param_value, "params": param_key}
-            else:
-                if "code" not in data:
-                    data["code"] = "unknown"
 
-            resp = self.make_response(data, status_code, headers)
+        # Remove duplicate Content-Length header
-        else:
+        remove_headers = ("Content-Length",)
-            if "code" not in data:
+        for header in remove_headers:
-                data["code"] = "unknown"
+            headers.pop(header, None)
 
-            resp = self.make_response(data, status_code, headers)
+        return api.make_response(data, status_code, headers)
 
-        if status_code == 401:
+
-            resp = self.unauthorized(resp)
+class ExternalApi(Api):
-        return resp
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        register_external_error_handlers(self)

api/mypy.ini

@@ -15,5 +15,8 @@ ignore_missing_imports=True
 [mypy-flask_restx]
 ignore_missing_imports=True
 
+[mypy-flask_restx.api]
+ignore_missing_imports=True
+
 [mypy-flask_restx.inputs]
 ignore_missing_imports=True