r2

api/controllers/console/datasets/rag_pipeline/datasource_auth.py

@@ -0,0 +1,140 @@
+
+from flask import redirect, request
+from flask_login import current_user  # type: ignore
+from flask_restful import (  # type: ignore
+    Resource,  # type: ignore
+    marshal_with,
+    reqparse,
+)
+from sqlalchemy.orm import Session
+from werkzeug.exceptions import Forbidden, NotFound
+
+from configs import dify_config
+from controllers.console import api
+from controllers.console.wraps import (
+    account_initialization_required,
+    setup_required,
+)
+from core.model_runtime.errors.validate import CredentialsValidateFailedError
+from core.plugin.impl.datasource import PluginDatasourceManager
+from core.plugin.impl.oauth import OAuthHandler
+from extensions.ext_database import db
+from libs.login import login_required
+from models.oauth import DatasourceOauthParamConfig, DatasourceProvider
+from services.datasource_provider_service import DatasourceProviderService
+
+
+class DatasourcePluginOauthApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider, plugin_id):
+        # Check user role first
+        if not current_user.is_editor:
+            raise Forbidden()
+        # get all plugin oauth configs
+        plugin_oauth_config = db.session.query(DatasourceOauthParamConfig).filter_by(
+            provider=provider,
+            plugin_id=plugin_id
+            ).first()
+        if not plugin_oauth_config:
+            raise NotFound()
+        oauth_handler = OAuthHandler()
+        redirect_url = f"{dify_config.CONSOLE_WEB_URL}/oauth/datasource/provider/{provider}/plugin/{plugin_id}/callback"
+        system_credentials = plugin_oauth_config.system_credentials
+        if system_credentials:
+            system_credentials["redirect_url"] = redirect_url
+        response = oauth_handler.get_authorization_url(
+            current_user.current_tenant.id,
+            current_user.id,
+            plugin_id,
+            provider,
+            system_credentials=system_credentials
+        )
+        return response.model_dump()
+
+class DatasourceOauthCallback(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, provider, plugin_id):
+        oauth_handler = OAuthHandler()
+        plugin_oauth_config = db.session.query(DatasourceOauthParamConfig).filter_by(
+            provider=provider,
+            plugin_id=plugin_id
+        ).first()
+        if not plugin_oauth_config:
+            raise NotFound()
+        credentials = oauth_handler.get_credentials(
+            current_user.current_tenant.id,
+            current_user.id,
+            plugin_id,
+            provider,
+            system_credentials=plugin_oauth_config.system_credentials,
+            request=request
+        )
+        datasource_provider = DatasourceProvider(
+            plugin_id=plugin_id,
+            provider=provider,
+            auth_type="oauth",
+            encrypted_credentials=credentials
+        )
+        db.session.add(datasource_provider)
+        db.session.commit()
+        return redirect(f"{dify_config.CONSOLE_WEB_URL}")
+
+class DatasourceAuth(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self, provider, plugin_id):
+        if not current_user.is_editor:
+            raise Forbidden()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")
+        args = parser.parse_args()
+
+        datasource_provider_service = DatasourceProviderService()
+
+        try:
+            datasource_provider_service.datasource_provider_credentials_validate(
+                tenant_id=current_user.current_tenant_id, 
+                provider=provider, 
+                plugin_id=plugin_id, 
+                credentials=args["credentials"]
+            )
+        except CredentialsValidateFailedError as ex:
+            raise ValueError(str(ex))
+
+        return {"result": "success"}, 201
+        
+class DatasourceAuthDeleteApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, provider, plugin_id):
+        if not current_user.is_editor:
+            raise Forbidden()
+        datasource_provider_service = DatasourceProviderService()
+        datasource_provider_service.remove_datasource_credentials(
+            tenant_id=current_user.current_tenant_id, 
+            provider=provider, 
+            plugin_id=plugin_id
+        )
+        return {"result": "success"}, 200
+
+# Import Rag Pipeline
+api.add_resource(
+    DatasourcePluginOauthApi,
+    "/oauth/datasource/provider/<string:provider>/plugin/<string:plugin_id>",
+)
+api.add_resource(
+    DatasourceOauthCallback,
+    "/oauth/datasource/provider/<string:provider>/plugin/<string:plugin_id>/callback",
+)
+api.add_resource(
+    DatasourceAuth,
+    "/auth/datasource/provider/<string:provider>/plugin/<string:plugin_id>",
+)
+

api/controllers/console/datasets/rag_pipeline/datasource_oauth.py

@@ -1,47 +0,0 @@
-
-from flask_login import current_user  # type: ignore
-from flask_restful import Resource  # type: ignore
-from werkzeug.exceptions import Forbidden
-
-from controllers.console import api
-from controllers.console.wraps import (
-    account_initialization_required,
-    setup_required,
-)
-from core.plugin.impl.datasource import PluginDatasourceManager
-from libs.login import login_required
-
-
-class DatasourcePluginOauthApi(Resource):
-    @setup_required
-    @login_required
-    @account_initialization_required
-    def get(self, datasource_type, datasource_name):
-        # Check user role first
-        if not current_user.is_editor:
-            raise Forbidden()
-        # get all builtin providers
-        manager = PluginDatasourceManager()
-        providers = manager.get_provider_oauth_url()
-        return providers
-
-
-
-
-# Import Rag Pipeline
-api.add_resource(
-    DatasourcePluginOauthApi,
-    "/datasource/<string:datasoruce_type>/<string:datasource_name>/oauth",
-)
-api.add_resource(
-    RagPipelineImportConfirmApi,
-    "/rag/pipelines/imports/<string:import_id>/confirm",
-)
-api.add_resource(
-    RagPipelineImportCheckDependenciesApi,
-    "/rag/pipelines/imports/<string:pipeline_id>/check-dependencies",
-)
-api.add_resource(
-    RagPipelineExportApi,
-    "/rag/pipelines/<string:pipeline_id>/exports",
-)