klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: assistant connect

Browse Source
This commit is contained in:
junjun
2025-07-29 11:47:12 +08:00
parent 17fcfd4503
commit 8141700746
2 changed files with 57 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+1 -1
View File
@@ -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
+56 -30
View File
@@ -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 @@ 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