diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 16736aa9..ae2eb9f9 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -881,7 +881,7 @@ class LLMService: def finish(self): return finish_record(session=self.session, record_id=self.record.id) - def execute_sql(self, sql: str, tables): + def execute_sql(self, sql: str): """Execute SQL query Args: @@ -893,7 +893,7 @@ class LLMService: """ SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}") try: - return exec_sql(ds=self.ds, sql=sql, origin_column=False, table_name=tables) + return exec_sql(ds=self.ds, sql=sql, origin_column=False) except Exception as e: if isinstance(e, ParseSQLResultError): raise e @@ -1022,7 +1022,6 @@ class LLMService: sql = self.check_save_sql(res=full_sql_text) else: sql = self.check_save_sql(res=full_sql_text) - tables = [] SQLBotLogUtil.info(sql) format_sql = sqlparse.format(sql, reindent=True) @@ -1040,7 +1039,7 @@ class LLMService: subsql) real_execute_sql = assistant_dynamic_sql - result = self.execute_sql(sql=real_execute_sql, tables=tables) + result = self.execute_sql(sql=real_execute_sql) self.save_sql_data(data_obj=result) if in_chat: yield 'data:' + orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n' diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index b4720bad..df55e3ab 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -297,7 +297,7 @@ def preview(session: SessionDep, current_user: CurrentUser, id: int, data: Table sql = f"""SELECT "{'", "'.join(fields)}" FROM "{data.table.table_name}" {where} LIMIT 100""" - return exec_sql(ds, sql, True, [data.table.table_name]) + return exec_sql(ds, sql, True) def fieldEnum(session: SessionDep, id: int): @@ -313,7 +313,7 @@ def fieldEnum(session: SessionDep, id: int): db = DB.get_db(ds.type) sql = f"""SELECT DISTINCT {db.prefix}{field.field_name}{db.suffix} FROM {db.prefix}{table.table_name}{db.suffix}""" - res = exec_sql(ds, sql, True, [table.table_name]) + res = exec_sql(ds, sql, True) return [item.get(res.get('fields')[0]) for item in res.get('data')] diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index a1da1492..5018c994 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -24,7 +24,7 @@ from apps.system.schemas.system_schema import AssistantOutDsSchema from common.core.deps import Trans from common.utils.utils import SQLBotLogUtil from fastapi import HTTPException -from apps.db.es_engine import get_es_connect, get_es_index, get_es_fields, get_es_data +from apps.db.es_engine import get_es_connect, get_es_index, get_es_fields, get_es_data_by_http def get_uri(ds: CoreDatasource) -> str: @@ -341,7 +341,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None): return res_list -def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False, table_name=None): +def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False): while sql.endswith(';'): sql = sql[:-1] @@ -421,17 +421,16 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column= raise ParseSQLResultError(str(ex)) elif ds.type == 'es': try: - if table_name and table_name[0]: - res, columns = get_es_data(conf, sql, table_name[0]) - columns = [field[0] for field in columns] if origin_column else [field[0].lower() for - field in - columns] - result_list = [ - {str(columns[i]): float(value) if isinstance(value, Decimal) else value for i, value in - enumerate(tuple_item)} - for tuple_item in res - ] - return {"fields": columns, "data": result_list, - "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))} + res, columns = get_es_data_by_http(conf, sql) + columns = [field.get('name') for field in columns] if origin_column else [field.get('name').lower() for + field in + columns] + result_list = [ + {str(columns[i]): float(value) if isinstance(value, Decimal) else value for i, value in + enumerate(tuple(tuple_item))} + for tuple_item in res + ] + return {"fields": columns, "data": result_list, + "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))} except Exception as ex: raise Exception(str(ex)) diff --git a/backend/apps/db/es_engine.py b/backend/apps/db/es_engine.py index 05c5797b..1353665d 100644 --- a/backend/apps/db/es_engine.py +++ b/backend/apps/db/es_engine.py @@ -2,10 +2,10 @@ # Date: 2025/9/9 import json +from base64 import b64encode import requests from elasticsearch import Elasticsearch -from fastapi import HTTPException from apps.datasource.models.datasource import DatasourceConf @@ -60,29 +60,55 @@ def get_es_fields(conf: DatasourceConf, table_name: str): return res -def get_es_data(conf: DatasourceConf, sql: str, table_name: str): - r = requests.post(f"{conf.host}/_sql/translate", json={"query": sql}) - if r.json().get('error'): - print(json.dumps(r.json())) +# def get_es_data(conf: DatasourceConf, sql: str, table_name: str): +# r = requests.post(f"{conf.host}/_sql/translate", json={"query": sql}) +# if r.json().get('error'): +# print(json.dumps(r.json())) +# +# es_client = get_es_connect(conf) +# response = es_client.search( +# index=table_name, +# body=json.dumps(r.json()) +# ) +# +# # print(response) +# fields = get_es_fields(conf, table_name) +# res = [] +# for hit in response.get('hits').get('hits'): +# item = [] +# if 'fields' in hit: +# result = hit.get('fields') # {'title': ['Python'], 'age': [30]} +# for field in fields: +# v = result.get(field[0]) +# item.append(v[0]) if v else item.append(None) +# res.append(tuple(item)) +# # print(hit['fields']['title'][0]) +# # elif '_source' in hit: +# # print(hit.get('_source')) +# return res, fields - es_client = get_es_connect(conf) - response = es_client.search( - index=table_name, - body=json.dumps(r.json()) - ) - # print(response) - fields = get_es_fields(conf, table_name) - res = [] - for hit in response.get('hits').get('hits'): - item = [] - if 'fields' in hit: - result = hit.get('fields') # {'title': ['Python'], 'age': [30]} - for field in fields: - v = result.get(field[0]) - item.append(v[0]) if v else item.append(None) - res.append(tuple(item)) - # print(hit['fields']['title'][0]) - # elif '_source' in hit: - # print(hit.get('_source')) - return res, fields +def get_es_data_by_http(conf: DatasourceConf, sql: str): + url = conf.host + while url.endswith('/'): + url = sql[:-1] + + host = f'{url}/_sql?format=json' + username = f"{conf.username}" + password = f"{conf.password}" + + credentials = f"{username}:{password}" + encoded_credentials = b64encode(credentials.encode()).decode() + + headers = { + "Content-Type": "application/json", + "Authorization": f"Basic {encoded_credentials}" + } + + response = requests.post(host, data=json.dumps({"query": sql}), headers=headers) + + # print(response.json()) + res = response.json() + fields = res.get('columns') + result = res.get('rows') + return result, fields diff --git a/frontend/src/views/ds/js/ds-type.ts b/frontend/src/views/ds/js/ds-type.ts index 7ff52a83..beeb7f43 100644 --- a/frontend/src/views/ds/js/ds-type.ts +++ b/frontend/src/views/ds/js/ds-type.ts @@ -7,7 +7,7 @@ import ck from '@/assets/datasource/icon_ck.png' import dm from '@/assets/datasource/icon_dm.png' import doris from '@/assets/datasource/icon_doris.png' import redshift from '@/assets/datasource/icon_redshift.png' -// import es from '@/assets/datasource/icon_es.png' +import es from '@/assets/datasource/icon_es.png' import { i18n } from '@/i18n' const t = i18n.global.t @@ -21,7 +21,7 @@ export const dsType = [ { label: '达梦', value: 'dm' }, { label: 'Apache Doris', value: 'doris' }, { label: 'AWS Redshift', value: 'redshift' }, - // { label: 'Elasticsearch', value: 'es' }, + { label: 'Elasticsearch', value: 'es' }, ] export const dsTypeWithImg = [ @@ -34,7 +34,7 @@ export const dsTypeWithImg = [ { name: '达梦', type: 'dm', img: dm }, { name: 'Apache Doris', type: 'doris', img: doris }, { name: 'AWS Redshift', type: 'redshift', img: redshift }, - // { name: 'Elasticsearch', type: 'es', img: es }, + { name: 'Elasticsearch', type: 'es', img: es }, ] export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm', 'redshift']