feat: support Elasticsearch datasource #108

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -625,7 +625,7 @@ def generate_assistant_dynamic_sql(self, sql, tables: List):
         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 @@ def save_sql_data(self, data_obj: Dict[str, Any]):
     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 @@ def execute_sql(self, sql: str):
         """
         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 @@ def run_task(self, in_chat: bool = True):
             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 @@ def run_task(self, in_chat: bool = True):
                     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 @@ def run_task(self, in_chat: bool = True):
             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

@@ -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

@@ -24,6 +24,7 @@
 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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -256,7 +256,8 @@
         "success": "Connect success",
         "failed": "Connect failed"
       },
-      "timeout": "Timeout(second)"
+      "timeout": "Timeout(second)",
+      "address": "Address"
     }
   },
   "datasource": {

frontend/src/i18n/en.json

@@ -256,7 +256,8 @@
         "success": "连接成功",
         "failed": "连接失败"
       },
-      "timeout": "查询超时(秒)"
+      "timeout": "查询超时(秒)",
+      "address": "地址"
     }
   },
   "datasource": {