Underground-Digital/Workflow-Engine

1from flask_login import current_user2from flask_restful import Resource, reqparse3from werkzeug.exceptions import Forbidden4 5from controllers.console import api6from controllers.console.auth.error import ApiKeyAuthFailedError7from libs.login import login_required8from services.auth.api_key_auth_service import ApiKeyAuthService9 10from ..wraps import account_initialization_required, setup_required11 12 13class ApiKeyAuthDataSource(Resource):14    @setup_required15    @login_required16    @account_initialization_required17    def get(self):18        data_source_api_key_bindings = ApiKeyAuthService.get_provider_auth_list(current_user.current_tenant_id)19        if data_source_api_key_bindings:20            return {21                "sources": [22                    {23                        "id": data_source_api_key_binding.id,24                        "category": data_source_api_key_binding.category,25                        "provider": data_source_api_key_binding.provider,26                        "disabled": data_source_api_key_binding.disabled,27                        "created_at": int(data_source_api_key_binding.created_at.timestamp()),28                        "updated_at": int(data_source_api_key_binding.updated_at.timestamp()),29                    }30                    for data_source_api_key_binding in data_source_api_key_bindings31                ]32            }33        return {"sources": []}34 35 36class ApiKeyAuthDataSourceBinding(Resource):37    @setup_required38    @login_required39    @account_initialization_required40    def post(self):41        # The role of the current user in the table must be admin or owner42        if not current_user.is_admin_or_owner:43            raise Forbidden()44        parser = reqparse.RequestParser()45        parser.add_argument("category", type=str, required=True, nullable=False, location="json")46        parser.add_argument("provider", type=str, required=True, nullable=False, location="json")47        parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json")48        args = parser.parse_args()49        ApiKeyAuthService.validate_api_key_auth_args(args)50        try:51            ApiKeyAuthService.create_provider_auth(current_user.current_tenant_id, args)52        except Exception as e:53            raise ApiKeyAuthFailedError(str(e))54        return {"result": "success"}, 20055 56 57class ApiKeyAuthDataSourceBindingDelete(Resource):58    @setup_required59    @login_required60    @account_initialization_required61    def delete(self, binding_id):62        # The role of the current user in the table must be admin or owner63        if not current_user.is_admin_or_owner:64            raise Forbidden()65 66        ApiKeyAuthService.delete_provider_auth(current_user.current_tenant_id, binding_id)67 68        return {"result": "success"}, 20069 70 71api.add_resource(ApiKeyAuthDataSource, "/api-key-auth/data-source")72api.add_resource(ApiKeyAuthDataSourceBinding, "/api-key-auth/data-source/binding")73api.add_resource(ApiKeyAuthDataSourceBindingDelete, "/api-key-auth/data-source/<uuid:binding_id>")74