diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 584e1dca..aa4eff58 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -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' diff --git a/backend/apps/db/constant.py b/backend/apps/db/constant.py index c4de1cd0..67074ac1 100644 --- a/backend/apps/db/constant.py +++ b/backend/apps/db/constant.py @@ -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 diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index fe85de6c..41cf503a 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -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')))} diff --git a/backend/apps/db/es_engine.py b/backend/apps/db/es_engine.py new file mode 100644 index 00000000..897251ed --- /dev/null +++ b/backend/apps/db/es_engine.py @@ -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 diff --git a/backend/apps/db/type.py b/backend/apps/db/type.py index e2c6c4fb..1e48dc65 100644 --- a/backend/apps/db/type.py +++ b/backend/apps/db/type.py @@ -13,5 +13,6 @@ def db_type_relation() -> Dict: "ck": "ClickHouse", "dm": "达梦", "doris": "Apache Doris", - "redshift": "AWS Redshift" + "redshift": "AWS Redshift", + "es": "Elasticsearch" } diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 6da8e3eb..a1c5b4e3 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -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] diff --git a/backend/template.yaml b/backend/template.yaml index dd182300..49f07c1e 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -70,7 +70,7 @@ template: 生成SQL时,必须避免与数据库关键字冲突 - 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦(DM)、AWS Redshift,则在schema、表名、字段名、别名外层加双引号; + 如数据库引擎是 PostgreSQL、Oracle、ClickHouse、达梦(DM)、AWS Redshift、Elasticsearch,则在schema、表名、字段名、别名外层加双引号; 如数据库引擎是 MySQL、Doris,则在表名、字段名、别名外层加反引号; 如数据库引擎是 Microsoft SQL Server,则在schema、表名、字段名、别名外层加方括号。 @@ -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格式返回: diff --git a/frontend/src/assets/datasource/icon_es.png b/frontend/src/assets/datasource/icon_es.png new file mode 100644 index 00000000..5d8f6e6e Binary files /dev/null and b/frontend/src/assets/datasource/icon_es.png differ diff --git a/frontend/src/i18n/en.json b/frontend/src/i18n/en.json index c0eef984..ea6d621c 100644 --- a/frontend/src/i18n/en.json +++ b/frontend/src/i18n/en.json @@ -256,7 +256,8 @@ "success": "Connect success", "failed": "Connect failed" }, - "timeout": "Timeout(second)" + "timeout": "Timeout(second)", + "address": "Address" } }, "datasource": { diff --git a/frontend/src/i18n/zh-CN.json b/frontend/src/i18n/zh-CN.json index 4cadc4c2..192f6827 100644 --- a/frontend/src/i18n/zh-CN.json +++ b/frontend/src/i18n/zh-CN.json @@ -256,7 +256,8 @@ "success": "连接成功", "failed": "连接失败" }, - "timeout": "查询超时(秒)" + "timeout": "查询超时(秒)", + "address": "地址" } }, "datasource": { diff --git a/frontend/src/views/ds/DatasourceForm.vue b/frontend/src/views/ds/DatasourceForm.vue index 47d9919a..72e82340 100644 --- a/frontend/src/views/ds/DatasourceForm.vue +++ b/frontend/src/views/ds/DatasourceForm.vue @@ -513,6 +513,7 @@ defineExpose({ 12+ 5.6+ 9.6+ + 7+
@@ -594,14 +595,21 @@ defineExpose({ />
- + - + - + {{ t('ds.form.mode.sid') }} - + - +