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: support Elasticsearch datasource #108

Browse Source
This commit is contained in:
junjun
2025-09-17 10:23:01 +08:00
parent f86b2f60cf
commit a0d603ca44
5 changed files with 72 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+3 -4
View File
@@ -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'
backend/apps/datasource/crud/datasource.py
+2 -2
View File
@@ -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')]
backend/apps/db/db.py
+13 -14
View File
@@ -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))
backend/apps/db/es_engine.py
+51 -25
View File
@@ -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
frontend/src/views/ds/js/ds-type.ts
+3 -3
View File
@@ -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']