feat: assistant connect

Author: XiaJunjie2020

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")

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="[email protected]")
+    return UserInfoDTO(id=id, account="sqlbot-inner-assistant", oid=1, name="sqlbot-inner-assistant",
+                       email="[email protected]")
+
 
 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 @@ def get_ds_from_api(self):
                     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 @@ def get_db_schema(self, ds_id: int) -> str:
                 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 @@ def get_db_schema(self, ds_id: int) -> str:
             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 @@ def convert2schema(self, ds_dict: dict) -> AssistantOutDsSchema:
         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
-