From 51e511a68a5cb3f0fdeff7c2e064649e9197c504 Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Tue, 16 Sep 2025 15:35:11 +0800
Subject: [PATCH] feat: support Elasticsearch datasource #108

---
 backend/apps/chat/task/llm.py | 34 ++++++++++++++++------------------
 backend/apps/db/db.py         |  2 --
 2 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py
index d8a422ea..30ac38a6 100644
--- a/backend/apps/chat/task/llm.py
+++ b/backend/apps/chat/task/llm.py
@@ -1000,26 +1000,26 @@ 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) 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)
-                else:
-                    sql_result = self.generate_filter(sql, tables)  # maybe no sql and tables
+                if is_normal_user(self.current_user):
+                    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)
+                    else:
+                        sql_result = self.generate_filter(sql, tables)  # maybe no sql and tables
 
-                if sql_result:
-                    SQLBotLogUtil.info(sql_result)
-                    sql = self.check_save_sql(res=sql_result)
-                elif dynamic_sql_result and sqlbot_temp_sql_text:
-                    assistant_dynamic_sql = self.check_save_sql(res=sqlbot_temp_sql_text)
-                else:
-                    sql = self.check_save_sql(res=full_sql_text)
+                    if sql_result:
+                        SQLBotLogUtil.info(sql_result)
+                        sql = self.check_save_sql(res=sql_result)
+                    elif dynamic_sql_result and sqlbot_temp_sql_text:
+                        assistant_dynamic_sql = self.check_save_sql(res=sqlbot_temp_sql_text)
+                    else:
+                        sql = self.check_save_sql(res=full_sql_text)
             else:
                 sql = self.check_save_sql(res=full_sql_text)
                 tables = []
@@ -1040,9 +1040,7 @@ class LLMService:
                                                                           subsql)
                 real_execute_sql = assistant_dynamic_sql
 
-            print(tables)
             result = self.execute_sql(sql=real_execute_sql, tables=tables)
-            print(result)
             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/db.py b/backend/apps/db/db.py
index 3a4dc94f..a1da1492 100644
--- a/backend/apps/db/db.py
+++ b/backend/apps/db/db.py
@@ -420,7 +420,6 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                 except Exception as ex:
                     raise ParseSQLResultError(str(ex))
         elif ds.type == 'es':
-            print(table_name)
             try:
                 if table_name and table_name[0]:
                     res, columns = get_es_data(conf, sql, table_name[0])
@@ -432,7 +431,6 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                          enumerate(tuple_item)}
                         for tuple_item in res
                     ]
-                    print(result_list)
                     return {"fields": columns, "data": result_list,
                             "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))}
             except Exception as ex: