Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/chat/task/llm.py

@@ -1000,26 +1000,26 @@ 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) 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 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 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)
             else:
                 sql = self.check_save_sql(res=full_sql_text)
                 tables = []

backend/apps/db/db.py

@@ -420,15 +420,18 @@ def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str, origin_column=
                 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')))}
+            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')))}
+            except Exception as ex:
+                raise Exception(str(ex))

backend/apps/db/es_engine.py

@@ -63,7 +63,7 @@ def get_es_fields(conf: DatasourceConf, table_name: str):
 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'):
-        raise HTTPException(status_code=500, detail=f': {json.dumps(r.json())}')
+        print(json.dumps(r.json()))
 
     es_client = get_es_connect(conf)
     response = es_client.search(