From 8141700746aa365cea14c56320b980efededc9c4 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 29 Jul 2025 11:47:12 +0800 Subject: [PATCH] feat: assistant connect --- backend/apps/datasource/crud/datasource.py | 2 +- backend/apps/system/crud/assistant.py | 86 ++++++++++++++-------- 2 files changed, 57 insertions(+), 31 deletions(-) diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index a73a22bf..ed52dfa4 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -41,7 +41,7 @@ def get_ds(session: SessionDep, id: int): def check_status(session: SessionDep, ds: CoreDatasource, is_raise: bool = False): - conn = get_engine(ds, 5) + conn = get_engine(ds, 10) try: with conn.connect() as connection: SQLBotLogUtil.info("success") diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 9b1579fc..72bdfa57 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -1,29 +1,33 @@ - - import json +import urllib from typing import Optional -from fastapi import FastAPI + import requests +from fastapi import FastAPI from sqlalchemy import Engine, create_engine from sqlmodel import Session, select -from apps.datasource.models.datasource import CoreDatasource, DatasourceConf +from starlette.middleware.cors import CORSMiddleware +from apps.datasource.models.datasource import CoreDatasource, DatasourceConf from apps.system.models.system_model import AssistantModel from apps.system.schemas.auth import CacheName, CacheNamespace from apps.system.schemas.system_schema import AssistantHeader, AssistantOutDsSchema, UserInfoDTO -from common.core.sqlbot_cache import cache -from common.core.db import engine -from starlette.middleware.cors import CORSMiddleware from common.core.config import settings +from common.core.db import engine +from common.core.sqlbot_cache import cache from common.utils.utils import string_to_numeric_hash + @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id") async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None: db_model = session.get(AssistantModel, assistant_id) return db_model + def get_assistant_user(*, id: int): - return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", email="sqlbot-inner-assistant@sqlbot.com") + return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant", + email="sqlbot-inner-assistant@sqlbot.com") + def get_assistant_ds(llm_service) -> list[dict]: assistant: AssistantHeader = llm_service.current_assistant @@ -34,13 +38,14 @@ def get_assistant_ds(llm_service) -> list[dict]: if configuration: config: dict[any] = json.loads(configuration) oid: int = int(config['oid']) - stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid) + stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where( + CoreDatasource.oid == oid) if not assistant.online: - private_list:list[int] = config.get('private_list') or None + private_list: list[int] = config.get('private_list') or None if private_list: stmt = stmt.where(~CoreDatasource.id.in_(private_list)) db_ds_list = session.exec(stmt) - + result_list = [ { "id": ds.id, @@ -49,7 +54,7 @@ def get_assistant_ds(llm_service) -> list[dict]: } for ds in db_ds_list ] - + # filter private ds if offline return result_list out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant) @@ -58,8 +63,9 @@ def get_assistant_ds(llm_service) -> list[dict]: # format? return dslist + def init_dynamic_cors(app: FastAPI): - try: + try: with Session(engine) as session: list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all() seen = set() @@ -81,20 +87,20 @@ def init_dynamic_cors(app: FastAPI): cors_middleware.kwargs['allow_origins'] = updated_origins except Exception as e: return False, e - - + class AssistantOutDs: assistant: AssistantHeader ds_list: Optional[list[AssistantOutDsSchema]] = None certificate: Optional[str] = None + def __init__(self, assistant: AssistantHeader): self.assistant = assistant self.ds_list = None self.certificate = assistant.certificate self.get_ds_from_api() - - #@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id") + + # @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id") def get_ds_from_api(self): config: dict[any] = json.loads(self.assistant.configuration) endpoint: str = config['endpoint'] @@ -118,23 +124,23 @@ class AssistantOutDs: self.convert2schema(item) for item in temp_list ] - + return self.ds_list else: raise Exception(f"Failed to get datasource list from {endpoint}, error: {result_json.get('message')}") else: raise Exception(f"Failed to get datasource list from {endpoint}, status code: {res.status_code}") - + def get_simple_ds_list(self): if self.ds_list: return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list] else: raise Exception("Datasource list is not found.") - + def get_db_schema(self, ds_id: int) -> str: ds = self.get_ds(ds_id) schema_str = "" - #db_name = ds.db_schema + # db_name = ds.db_schema db_name = ds.db_schema if ds.db_schema is not None and ds.db_schema != "" else ds.dataBase schema_str += f"【DB_ID】 {db_name}\n【Schema】\n" for table in ds.tables: @@ -144,7 +150,7 @@ class AssistantOutDs: schema_str += '\n[\n' else: schema_str += f", {table_comment}\n[\n" - + field_list = [] for field in table.fields: field_comment = field.comment @@ -155,7 +161,7 @@ class AssistantOutDs: schema_str += ",\n".join(field_list) schema_str += '\n]\n' return schema_str - + def get_ds(self, ds_id: int): if self.ds_list: for ds in self.ds_list: @@ -175,20 +181,22 @@ class AssistantOutDs: db_schema = ds_dict.get('schema', ds_dict.get('db_schema', '')) ds_dict.pop("schema", None) return AssistantOutDsSchema(**{**ds_dict, "id": id, "db_schema": db_schema}) - + + class AssistantOutDsFactory: @staticmethod def get_instance(assistant: AssistantHeader) -> AssistantOutDs: return AssistantOutDs(assistant) + def get_ds_engine(ds: AssistantOutDsSchema) -> Engine: timeout: int = 30 connect_args = {"connect_timeout": timeout} conf = DatasourceConf( - host=ds.host, - port=ds.port, + host=ds.host, + port=ds.port, username=ds.user, - password=ds.password, + password=ds.password, database=ds.dataBase, driver='', extraJdbc=ds.extraParams, @@ -197,8 +205,26 @@ def get_ds_engine(ds: AssistantOutDsSchema) -> Engine: conf.extraJdbc = '' from apps.db.db import get_uri_from_config uri = get_uri_from_config(ds.type, conf) + # if ds.type == "pg" and ds.db_schema: + # connect_args.update({"options": f"-c search_path={ds.db_schema}"}) + # engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10) + if ds.type == "pg" and ds.db_schema: - connect_args.update({"options": f"-c search_path={ds.db_schema}"}) - engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10) + engine = create_engine(uri, + connect_args={"options": f"-c search_path={urllib.parse.quote(ds.db_schema)}", + "connect_timeout": timeout}, + pool_timeout=timeout, pool_size=20, max_overflow=10) + elif ds.type == 'sqlServer': + engine = create_engine(uri, pool_timeout=timeout, + pool_size=20, + max_overflow=10) + elif ds.type == 'oracle': + engine = create_engine(uri, + pool_timeout=timeout, + pool_size=20, + max_overflow=10) + else: + engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout, + pool_size=20, + max_overflow=10) return engine - \ No newline at end of file