feat: support Elasticsearch datasource #108

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -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, tables):
+    def execute_sql(self, sql: str):
         """Execute SQL query
 
         Args:
@@ -893,7 +893,7 @@ def execute_sql(self, sql: str, tables):
         """
         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 @@ 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)
@@ -1040,7 +1039,7 @@ def run_task(self, in_chat: bool = True):
                                                                           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

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

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

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

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