refactor: assistant support all sqlbot datasource

This commit is contained in:
junjun
2025-12-04 17:35:01 +08:00
committed by Junjun
parent 39f1195a3a
commit d92c377e59
+18 -46
View File
@@ -117,7 +117,8 @@ class AssistantOutDs:
endpoint: str = config['endpoint']
endpoint = self.get_complete_endpoint(endpoint=endpoint)
if not endpoint:
raise Exception(f"Failed to get datasource list from {config['endpoint']}, error: [Assistant domain or endpoint miss]")
raise Exception(
f"Failed to get datasource list from {config['endpoint']}, error: [Assistant domain or endpoint miss]")
certificateList: list[any] = json.loads(self.certificate)
header = {}
cookies = {}
@@ -157,7 +158,9 @@ class AssistantOutDs:
if not domain_text:
return None
if ',' in domain_text or ';' in domain_text:
return (self.request_origin.strip('/') if self.request_origin else self.get_first_element(domain_text).strip('/')) + endpoint
return (
self.request_origin.strip('/') if self.request_origin else self.get_first_element(domain_text).strip(
'/')) + endpoint
else:
return f"{domain_text}{endpoint}"
@@ -206,14 +209,15 @@ class AssistantOutDs:
return schema_str
def get_ds(self, ds_id: int,trans: Trans = None):
def get_ds(self, ds_id: int, trans: Trans = None):
if self.ds_list:
for ds in self.ds_list:
if ds.id == ds_id:
return ds
else:
raise Exception("Datasource list is not found.")
raise Exception(f"Datasource id {ds_id} is not found." if trans is None else trans('i18n_data_training.datasource_id_not_found', key=ds_id))
raise Exception(f"Datasource id {ds_id} is not found." if trans is None else trans(
'i18n_data_training.datasource_id_not_found', key=ds_id))
def convert2schema(self, ds_dict: dict, config: dict[any]) -> AssistantOutDsSchema:
id_marker: str = ''
@@ -244,49 +248,17 @@ class AssistantOutDsFactory:
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,
username=ds.user,
password=ds.password,
database=ds.dataBase,
driver='',
extraJdbc=ds.extraParams or '',
dbSchema=ds.db_schema or ''
)
conf.extraJdbc = ''
from apps.db.db import get_uri_from_config
uri = get_uri_from_config(ds.type, conf)
if equals_ignore_case(ds.type, "pg") and ds.db_schema:
engine = create_engine(uri,
connect_args={"options": f"-c search_path={urllib.parse.quote(ds.db_schema)}",
"connect_timeout": timeout},
pool_timeout=timeout)
elif equals_ignore_case(ds.type, 'sqlServer'):
engine = create_engine(uri, pool_timeout=timeout)
elif equals_ignore_case(ds.type, 'oracle'):
engine = create_engine(uri,
pool_timeout=timeout)
else:
engine = create_engine(uri, connect_args={"connect_timeout": timeout}, pool_timeout=timeout)
return engine
def get_out_ds_conf(ds: AssistantOutDsSchema, timeout:int=30) -> str:
def get_out_ds_conf(ds: AssistantOutDsSchema, timeout: int = 30) -> str:
conf = {
"host":ds.host or '',
"port":ds.port or 0,
"username":ds.user or '',
"password":ds.password or '',
"database":ds.dataBase or '',
"driver":'',
"extraJdbc":ds.extraParams or '',
"dbSchema":ds.db_schema or '',
"timeout":timeout or 30
"host": ds.host or '',
"port": ds.port or 0,
"username": ds.user or '',
"password": ds.password or '',
"database": ds.dataBase or '',
"driver": '',
"extraJdbc": ds.extraParams or '',
"dbSchema": ds.db_schema or '',
"timeout": timeout or 30
}
conf["extraJdbc"] = ''
return aes_encrypt(json.dumps(conf))