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

feat: support Elasticsearch datasource #108

Browse Source
This commit is contained in:
junjun
2025-09-16 10:27:11 +08:00
parent 10129b785f
commit 0b588c4240
12 changed files with 163 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+14 -10
View File
@@ -625,7 +625,7 @@ class LLMService:
result_dict = {}
for table in ds.tables:
if table.name in tables and table.sql:
#sub_query.append({"table": table.name, "query": table.sql})
# sub_query.append({"table": table.name, "query": table.sql})
result_dict[table.name] = table.sql
sub_query.append({"table": table.name, "query": f'{dynamic_subsql_prefix}{table.name}'})
if not sub_query:
@@ -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):
def execute_sql(self, sql: str, tables):
"""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(self.ds, sql)
return exec_sql(ds=self.ds, sql=sql, origin_column=False, table_name=tables)
except Exception as e:
if isinstance(e, ParseSQLResultError):
raise e
@@ -1000,14 +1000,16 @@ class LLMService:
sqlbot_temp_sql_text = None
assistant_dynamic_sql = None
# todo row permission
if ((not self.current_assistant or is_page_embedded) and is_normal_user(self.current_user)) or use_dynamic_ds:
if ((not self.current_assistant or is_page_embedded) and is_normal_user(
self.current_user)) or use_dynamic_ds:
sql, tables = self.check_sql(res=full_sql_text)
sql_result = None
if use_dynamic_ds:
dynamic_sql_result = self.generate_assistant_dynamic_sql(sql, tables)
sqlbot_temp_sql_text = dynamic_sql_result.get('sqlbot_temp_sql_text') if dynamic_sql_result else None
#sql_result = self.generate_assistant_filter(sql, tables)
sqlbot_temp_sql_text = dynamic_sql_result.get(
'sqlbot_temp_sql_text') if dynamic_sql_result else None
# sql_result = self.generate_assistant_filter(sql, tables)
else:
sql_result = self.generate_filter(sql, tables) # maybe no sql and tables
@@ -1020,6 +1022,7 @@ 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)
@@ -1033,10 +1036,11 @@ class LLMService:
if sqlbot_temp_sql_text and assistant_dynamic_sql:
dynamic_sql_result.pop('sqlbot_temp_sql_text')
for origin_table, subsql in dynamic_sql_result.items():
assistant_dynamic_sql = assistant_dynamic_sql.replace(f'{dynamic_subsql_prefix}{origin_table}', subsql)
assistant_dynamic_sql = assistant_dynamic_sql.replace(f'{dynamic_subsql_prefix}{origin_table}',
subsql)
real_execute_sql = assistant_dynamic_sql
result = self.execute_sql(sql=real_execute_sql)
result = self.execute_sql(sql=real_execute_sql, tables=tables)
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/db/constant.py
+1
View File
@@ -22,6 +22,7 @@ class DB(Enum):
dm = ('dm', '"', '"', ConnectType.py_driver)
doris = ('doris', '`', '`', ConnectType.py_driver)
redshift = ('redshift', '"', '"', ConnectType.py_driver)
es = ('es', '"', '"', ConnectType.py_driver)
def __init__(self, type, prefix, suffix, connect_type: ConnectType):
self.type = type
backend/apps/db/db.py
+34 -3
View File
@@ -24,6 +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
def get_uri(ds: CoreDatasource) -> str:
@@ -144,7 +145,8 @@ def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDs
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
return False
elif ds.type == 'redshift':
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database, user=conf.username,
with redshift_connector.connect(host=conf.host, port=conf.port, database=conf.database,
user=conf.username,
password=conf.password,
timeout=10) as conn, conn.cursor() as cursor:
try:
@@ -156,6 +158,14 @@ def check_connection(trans: Optional[Trans], ds: CoreDatasource | AssistantOutDs
if is_raise:
raise HTTPException(status_code=500, detail=trans('i18n_ds_invalid') + f': {e.args}')
return False
elif ds.type == 'es':
es_conn = get_es_connect(conf)
if es_conn.ping():
SQLBotLogUtil.info("success")
return True
else:
SQLBotLogUtil.info("failed")
return False
else:
conn = get_ds_engine(ds)
try:
@@ -208,7 +218,7 @@ def get_version(ds: CoreDatasource | AssistantOutDsSchema):
cursor.execute(sql)
res = cursor.fetchall()
version = res[0][0]
elif ds.type == 'redshift':
elif ds.type == 'redshift' or ds.type == 'es':
version = ''
except Exception as e:
print(e)
@@ -285,6 +295,10 @@ def get_tables(ds: CoreDatasource):
res = cursor.fetchall()
res_list = [TableSchema(*item) for item in res]
return res_list
elif ds.type == 'es':
res = get_es_index(conf)
res_list = [TableSchema(*item) for item in res]
return res_list
def get_fields(ds: CoreDatasource, table_name: str = None):
@@ -321,9 +335,13 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
res = cursor.fetchall()
res_list = [ColumnSchema(*item) for item in res]
return res_list
elif ds.type == 'es':
res = get_es_fields(conf, table_name)
res_list = [ColumnSchema(*item) for item in res]
return res_list
def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False):
def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=False, table_name=None):
while sql.endswith(';'):
sql = sql[:-1]
@@ -401,3 +419,16 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
"sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
except Exception as ex:
raise ParseSQLResultError(str(ex))
elif ds.type == 'es':
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')))}
backend/apps/db/es_engine.py
+84
View File
@@ -0,0 +1,84 @@
# Author: Junjun
# Date: 2025/9/9
import json
import requests
from elasticsearch import Elasticsearch
from apps.datasource.models.datasource import DatasourceConf
def get_es_connect(conf: DatasourceConf):
es_client = Elasticsearch(
[conf.host], # ES address
basic_auth=(conf.username, conf.password),
verify_certs=False,
compatibility_mode=True
)
return es_client
# get tables
def get_es_index(conf: DatasourceConf):
es_client = get_es_connect(conf)
indices = es_client.cat.indices(format="json")
res = []
for idx in indices:
index_name = idx.get('index')
desc = ''
# get mapping
mapping = es_client.indices.get_mapping(index=index_name)
mappings = mapping.get(index_name).get("mappings")
if mappings.get('_meta'):
desc = mappings.get('_meta').get('description')
res.append((index_name, desc))
return res
# get fields
def get_es_fields(conf: DatasourceConf, table_name: str):
es_client = get_es_connect(conf)
index_name = table_name
mapping = es_client.indices.get_mapping(index=index_name)
properties = mapping.get(index_name).get("mappings").get("properties")
res = []
for field, config in properties.items():
field_type = config.get("type")
desc = ''
if config.get("_meta"):
desc = config.get("_meta").get('description')
if field_type:
res.append((field, field_type, desc))
else:
# object、nested...
res.append((field, ','.join(list(config.keys())), desc))
return res
def get_es_data(conf: DatasourceConf, sql: str, table_name: str):
r = requests.post(f"{conf.host}/_sql/translate", json={"query": sql})
# 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
backend/apps/db/type.py
+2 -1
View File
@@ -13,5 +13,6 @@ def db_type_relation() -> Dict:
"ck": "ClickHouse",
"dm": "达梦",
"doris": "Apache Doris",
"redshift": "AWS Redshift"
"redshift": "AWS Redshift",
"es": "Elasticsearch"
}
backend/pyproject.toml
+1
View File
@@ -50,6 +50,7 @@ dependencies = [
"dicttoxml>=1.7.16",
"dmpython>=2.5.22; platform_system != 'Darwin'",
"redshift-connector>=2.1.8",
"elasticsearch[requests] (>=7.10,<8.0)",
]
[project.optional-dependencies]
backend/template.yaml
+2 -2
View File
@@ -70,7 +70,7 @@ template:
生成SQL时，必须避免与数据库关键字冲突
</rule>
<rule>
如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift，则在schema、表名、字段名、别名外层加双引号；
如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift、Elasticsearch，则在schema、表名、字段名、别名外层加双引号；
如数据库引擎是 MySQL、Doris，则在表名、字段名、别名外层加反引号；
如数据库引擎是 Microsoft SQL Server，则在schema、表名、字段名、别名外层加方括号。
<example>
@@ -448,7 +448,7 @@ template:
- 如果存在冗余的过滤条件则进行去重后再生成新SQL。
- 给过滤条件中的字段前加上表别名（如果没有表别名则加表名），如：table.field。
- 生成SQL时，必须避免关键字冲突:
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift，则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦（DM）、AWS Redshift、Elasticsearch，则在schema、表名、字段名、别名外层加双引号；
- 如数据库引擎是 MySQL、Doris，则在表名、字段名、别名外层加反引号；
- 如数据库引擎是 Microsoft SQL Server，则在schema、表名、字段名、别名外层加方括号。
- 生成的SQL使用JSON格式返回：
frontend/src/assets/datasource/icon_es.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.0 KiB

frontend/src/i18n/en.json
+2 -1
View File
@@ -256,7 +256,8 @@
"success": "Connect success",
"failed": "Connect failed"
},
"timeout": "Timeout(second)"
"timeout": "Timeout(second)",
"address": "Address"
}
},
"datasource": {
frontend/src/i18n/zh-CN.json
+2 -1
View File
@@ -256,7 +256,8 @@
"success": "连接成功",
"failed": "连接失败"
},
"timeout": "查询超时(秒)"
"timeout": "查询超时(秒)",
"address": "地址"
}
},
"datasource": {
frontend/src/views/ds/DatasourceForm.vue
+18 -6
View File
@@ -513,6 +513,7 @@ defineExpose({
<span v-else-if="form.type === 'oracle'">12+</span>
<span v-else-if="form.type === 'mysql'">5.6+</span>
<span v-else-if="form.type === 'pg'">9.6+</span>
<span v-else-if="form.type === 'es'">7+</span>
</span>
</div>
<div class="form-content">
@@ -594,14 +595,21 @@ defineExpose({
/>
</el-form-item>
<div v-if="form.type !== 'excel'" style="margin-top: 16px">
<el-form-item :label="t('ds.form.host')" prop="host">
<el-form-item
:label="form.type !== 'es' ? t('ds.form.host') : t('ds.form.address')"
prop="host"
>
<el-input
v-model="form.host"
clearable
:placeholder="$t('datasource.please_enter') + $t('common.empty') + t('ds.form.host')"
:placeholder="
$t('datasource.please_enter') +
$t('common.empty') +
(form.type !== 'es' ? t('ds.form.host') : t('ds.form.address'))
"
/>
</el-form-item>
<el-form-item :label="t('ds.form.port')" prop="port">
<el-form-item v-if="form.type !== 'es'" :label="t('ds.form.port')" prop="port">
<el-input
v-model="form.port"
clearable
@@ -628,7 +636,11 @@ defineExpose({
show-password
/>
</el-form-item>
<el-form-item v-if="form.type !== 'dm'" :label="t('ds.form.database')" prop="database">
<el-form-item
v-if="form.type !== 'dm' && form.type !== 'es'"
:label="t('ds.form.database')"
prop="database"
>
<el-input
v-model="form.database"
clearable
@@ -647,7 +659,7 @@ defineExpose({
<el-radio value="sid">{{ t('ds.form.mode.sid') }}</el-radio>
</el-radio-group>
</el-form-item>
<el-form-item :label="t('ds.form.extra_jdbc')">
<el-form-item v-if="form.type !== 'es'" :label="t('ds.form.extra_jdbc')">
<el-input
v-model="form.extraJdbc"
clearable
@@ -681,7 +693,7 @@ defineExpose({
/>
</el-select>
</el-form-item>
<el-form-item :label="t('ds.form.timeout')" prop="timeout">
<el-form-item v-if="form.type !== 'es'" :label="t('ds.form.timeout')" prop="timeout">
<el-input-number
v-model="form.timeout"
clearable
frontend/src/views/ds/js/ds-type.ts
+3
View File
@@ -7,6 +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 { i18n } from '@/i18n'
const t = i18n.global.t
@@ -20,6 +21,7 @@ export const dsType = [
{ label: '达梦', value: 'dm' },
{ label: 'Apache Doris', value: 'doris' },
{ label: 'AWS Redshift', value: 'redshift' },
{ label: 'Elasticsearch', value: 'es' },
]
export const dsTypeWithImg = [
@@ -32,6 +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 },
]
export const haveSchema = ['sqlServer', 'pg', 'oracle', 'dm', 'redshift']