Underground-Digital/Workflow-Engine

1import logging2 3import requests4from flask import current_app, redirect, request5from flask_login import current_user6from flask_restful import Resource7from werkzeug.exceptions import Forbidden8 9from configs import dify_config10from controllers.console import api11from libs.login import login_required12from libs.oauth_data_source import NotionOAuth13 14from ..wraps import account_initialization_required, setup_required15 16 17def get_oauth_providers():18    with current_app.app_context():19        notion_oauth = NotionOAuth(20            client_id=dify_config.NOTION_CLIENT_ID,21            client_secret=dify_config.NOTION_CLIENT_SECRET,22            redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion",23        )24 25        OAUTH_PROVIDERS = {"notion": notion_oauth}26        return OAUTH_PROVIDERS27 28 29class OAuthDataSource(Resource):30    def get(self, provider: str):31        # The role of the current user in the table must be admin or owner32        if not current_user.is_admin_or_owner:33            raise Forbidden()34        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()35        with current_app.app_context():36            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)37            print(vars(oauth_provider))38        if not oauth_provider:39            return {"error": "Invalid provider"}, 40040        if dify_config.NOTION_INTEGRATION_TYPE == "internal":41            internal_secret = dify_config.NOTION_INTERNAL_SECRET42            if not internal_secret:43                return ({"error": "Internal secret is not set"},)44            oauth_provider.save_internal_access_token(internal_secret)45            return {"data": ""}46        else:47            auth_url = oauth_provider.get_authorization_url()48            return {"data": auth_url}, 20049 50 51class OAuthDataSourceCallback(Resource):52    def get(self, provider: str):53        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()54        with current_app.app_context():55            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)56        if not oauth_provider:57            return {"error": "Invalid provider"}, 40058        if "code" in request.args:59            code = request.args.get("code")60 61            return redirect(f"{dify_config.CONSOLE_WEB_URL}?type=notion&code={code}")62        elif "error" in request.args:63            error = request.args.get("error")64 65            return redirect(f"{dify_config.CONSOLE_WEB_URL}?type=notion&error={error}")66        else:67            return redirect(f"{dify_config.CONSOLE_WEB_URL}?type=notion&error=Access denied")68 69 70class OAuthDataSourceBinding(Resource):71    def get(self, provider: str):72        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()73        with current_app.app_context():74            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)75        if not oauth_provider:76            return {"error": "Invalid provider"}, 40077        if "code" in request.args:78            code = request.args.get("code")79            try:80                oauth_provider.get_access_token(code)81            except requests.exceptions.HTTPError as e:82                logging.exception(83                    f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}"84                )85                return {"error": "OAuth data source process failed"}, 40086 87            return {"result": "success"}, 20088 89 90class OAuthDataSourceSync(Resource):91    @setup_required92    @login_required93    @account_initialization_required94    def get(self, provider, binding_id):95        provider = str(provider)96        binding_id = str(binding_id)97        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()98        with current_app.app_context():99            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)100        if not oauth_provider:101            return {"error": "Invalid provider"}, 400102        try:103            oauth_provider.sync_data_source(binding_id)104        except requests.exceptions.HTTPError as e:105            logging.exception(f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")106            return {"error": "OAuth data source process failed"}, 400107 108        return {"result": "success"}, 200109 110 111api.add_resource(OAuthDataSource, "/oauth/data-source/<string:provider>")112api.add_resource(OAuthDataSourceCallback, "/oauth/data-source/callback/<string:provider>")113api.add_resource(OAuthDataSourceBinding, "/oauth/data-source/binding/<string:provider>")114api.add_resource(OAuthDataSourceSync, "/oauth/data-source/<string:provider>/<uuid:binding_id>/sync")115