feat/datasource
parent
83ca7f8deb
commit
5fc2bc58a9
@ -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>",
|
||||||
|
)
|
||||||
|
|
||||||
@ -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",
|
|
||||||
)
|
|
||||||
@ -0,0 +1,150 @@
|
|||||||
|
import logging
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from flask_login import current_user
|
||||||
|
|
||||||
|
from constants import HIDDEN_VALUE
|
||||||
|
from core import datasource
|
||||||
|
from core.datasource.__base import datasource_provider
|
||||||
|
from core.entities.model_entities import ModelStatus, ModelWithProviderEntity, ProviderModelWithStatusEntity
|
||||||
|
from core.helper import encrypter
|
||||||
|
from core.model_runtime.entities.model_entities import ModelType, ParameterRule
|
||||||
|
from core.model_runtime.entities.provider_entities import FormType
|
||||||
|
from core.model_runtime.errors.validate import CredentialsValidateFailedError
|
||||||
|
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
|
||||||
|
from core.plugin.impl.datasource import PluginDatasourceManager
|
||||||
|
from core.provider_manager import ProviderManager
|
||||||
|
from models.oauth import DatasourceProvider
|
||||||
|
from models.provider import ProviderType
|
||||||
|
from services.entities.model_provider_entities import (
|
||||||
|
CustomConfigurationResponse,
|
||||||
|
CustomConfigurationStatus,
|
||||||
|
DefaultModelResponse,
|
||||||
|
ModelWithProviderEntityResponse,
|
||||||
|
ProviderResponse,
|
||||||
|
ProviderWithModelsResponse,
|
||||||
|
SimpleProviderEntityResponse,
|
||||||
|
SystemConfigurationResponse,
|
||||||
|
)
|
||||||
|
from extensions.database import db
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class DatasourceProviderService:
|
||||||
|
"""
|
||||||
|
Model Provider Service
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.provider_manager = PluginDatasourceManager()
|
||||||
|
|
||||||
|
def datasource_provider_credentials_validate(self,
|
||||||
|
tenant_id: str,
|
||||||
|
provider: str,
|
||||||
|
plugin_id: str,
|
||||||
|
credentials: dict) -> None:
|
||||||
|
"""
|
||||||
|
validate datasource provider credentials.
|
||||||
|
|
||||||
|
:param tenant_id:
|
||||||
|
:param provider:
|
||||||
|
:param credentials:
|
||||||
|
"""
|
||||||
|
credential_valid = self.provider_manager.validate_provider_credentials(tenant_id=tenant_id,
|
||||||
|
user_id=current_user.id,
|
||||||
|
provider=provider,
|
||||||
|
credentials=credentials)
|
||||||
|
if credential_valid:
|
||||||
|
# Get all provider configurations of the current workspace
|
||||||
|
datasource_provider = db.session.query(DatasourceProvider).filter_by(tenant_id=tenant_id,
|
||||||
|
provider=provider,
|
||||||
|
plugin_id=plugin_id).first()
|
||||||
|
|
||||||
|
provider_credential_secret_variables = self.extract_secret_variables(tenant_id=tenant_id,
|
||||||
|
provider=provider
|
||||||
|
)
|
||||||
|
if not datasource_provider:
|
||||||
|
for key, value in credentials.items():
|
||||||
|
if key in provider_credential_secret_variables:
|
||||||
|
# if send [__HIDDEN__] in secret input, it will be same as original value
|
||||||
|
credentials[key] = encrypter.encrypt_token(tenant_id, value)
|
||||||
|
datasource_provider = DatasourceProvider(tenant_id=tenant_id,
|
||||||
|
provider=provider,
|
||||||
|
plugin_id=plugin_id,
|
||||||
|
auth_type="api_key",
|
||||||
|
encrypted_credentials=credentials)
|
||||||
|
db.session.add(datasource_provider)
|
||||||
|
db.session.commit()
|
||||||
|
else:
|
||||||
|
original_credentials = datasource_provider.encrypted_credentials
|
||||||
|
for key, value in credentials.items():
|
||||||
|
if key in provider_credential_secret_variables:
|
||||||
|
# if send [__HIDDEN__] in secret input, it will be same as original value
|
||||||
|
if value == HIDDEN_VALUE and key in original_credentials:
|
||||||
|
original_value = encrypter.encrypt_token(tenant_id, original_credentials[key])
|
||||||
|
credentials[key] = encrypter.encrypt_token(tenant_id, original_value)
|
||||||
|
else:
|
||||||
|
credentials[key] = encrypter.encrypt_token(tenant_id, value)
|
||||||
|
|
||||||
|
datasource_provider.encrypted_credentials = credentials
|
||||||
|
db.session.commit()
|
||||||
|
else:
|
||||||
|
raise CredentialsValidateFailedError()
|
||||||
|
|
||||||
|
def extract_secret_variables(self, tenant_id: str, provider: str) -> list[str]:
|
||||||
|
"""
|
||||||
|
Extract secret input form variables.
|
||||||
|
|
||||||
|
:param credential_form_schemas:
|
||||||
|
:return:
|
||||||
|
"""
|
||||||
|
datasource_provider = self.provider_manager.fetch_datasource_provider(tenant_id=tenant_id, provider=provider)
|
||||||
|
credential_form_schemas = datasource_provider.declaration.credentials_schema
|
||||||
|
secret_input_form_variables = []
|
||||||
|
for credential_form_schema in credential_form_schemas:
|
||||||
|
if credential_form_schema.type == FormType.SECRET_INPUT:
|
||||||
|
secret_input_form_variables.append(credential_form_schema.name)
|
||||||
|
|
||||||
|
return secret_input_form_variables
|
||||||
|
|
||||||
|
|
||||||
|
def get_datasource_credentials(self, tenant_id: str, provider: str, plugin_id: str) -> Optional[dict]:
|
||||||
|
"""
|
||||||
|
get datasource credentials.
|
||||||
|
|
||||||
|
:param tenant_id: workspace id
|
||||||
|
:param provider: provider name
|
||||||
|
:param datasource_name: datasource name
|
||||||
|
:param plugin_id: plugin id
|
||||||
|
:return:
|
||||||
|
"""
|
||||||
|
# Get all provider configurations of the current workspace
|
||||||
|
datasource_provider = db.session.query(DatasourceProvider).filter_by(tenant_id=tenant_id,
|
||||||
|
provider=provider,
|
||||||
|
plugin_id=plugin_id).first()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def remove_datasource_credentials(self,
|
||||||
|
tenant_id: str,
|
||||||
|
provider: str,
|
||||||
|
plugin_id: str) -> None:
|
||||||
|
"""
|
||||||
|
remove datasource credentials.
|
||||||
|
|
||||||
|
:param tenant_id: workspace id
|
||||||
|
:param provider: provider name
|
||||||
|
:param datasource_name: datasource name
|
||||||
|
:param plugin_id: plugin id
|
||||||
|
:return:
|
||||||
|
"""
|
||||||
|
# Get all provider configurations of the current workspace
|
||||||
|
datasource_provider = db.session.query(DatasourceProvider).filter_by(tenant_id=tenant_id,
|
||||||
|
provider=provider,
|
||||||
|
plugin_id=plugin_id).first()
|
||||||
|
if datasource_provider:
|
||||||
|
db.session.delete(datasource_provider)
|
||||||
|
db.session.commit()
|
||||||
|
|
||||||
Loading…
Reference in New Issue