From f2c4d53c582c6c796fd65007eb7bee08bc5b680f Mon Sep 17 00:00:00 2001 From: H <43509927+guoyuhao2330@users.noreply.github.com> Date: Thu, 22 Aug 2024 10:19:22 +0800 Subject: [PATCH] Fix component exesql bug (#2042) ### What problem does this PR solve? ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --------- Co-authored-by: Kevin Hu --- agent/canvas.py | 2 +- agent/component/exesql.py | 7 +++++-- api/apps/canvas_app.py | 2 +- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/agent/canvas.py b/agent/canvas.py index 566a2bf31..05d6e31f3 100644 --- a/agent/canvas.py +++ b/agent/canvas.py @@ -274,7 +274,7 @@ class Canvas(ABC): def get_embedding_model(self): return self._embed_id - def _find_loop(self, max_loops=2): + def _find_loop(self, max_loops=6): path = self.path[-1][::-1] if len(path) < 2: return False diff --git a/agent/component/exesql.py b/agent/component/exesql.py index 79d61d3ae..a5ce09e74 100644 --- a/agent/component/exesql.py +++ b/agent/component/exesql.py @@ -14,7 +14,7 @@ # limitations under the License. # from abc import ABC - +import re import pandas as pd from peewee import MySQLDatabase, PostgresqlDatabase from agent.component.base import ComponentBase, ComponentParamBase @@ -59,6 +59,9 @@ class ExeSQL(ComponentBase, ABC): ans = self.get_input() ans = "".join(ans["content"]) if "content" in ans else "" + ans = re.sub(r'^.*?SELECT ', 'SELECT ', repr(ans), flags=re.IGNORECASE) + ans = re.sub(r';.*?SELECT ', '; SELECT ', ans, flags=re.IGNORECASE) + ans = re.sub(r';[^;]*$', r';', ans) if not ans: return ExeSQL.be_output("SQL statement not found!") @@ -75,7 +78,7 @@ class ExeSQL(ComponentBase, ABC): sql_res = [{"content": rec + "\n"} for rec in [str(i) for i in query.fetchall()]] db.close() except Exception as e: - return ExeSQL.be_output("**Error**:" + str(e)) + return ExeSQL.be_output("**Error**:" + str(e) + "\nError SQL Statement:" + ans) if not sql_res: return ExeSQL.be_output("No record in the database!") diff --git a/api/apps/canvas_app.py b/api/apps/canvas_app.py index 0e7470513..7f049e4fc 100644 --- a/api/apps/canvas_app.py +++ b/api/apps/canvas_app.py @@ -175,6 +175,6 @@ def test_db_connect(): password=req["password"]) db.connect() db.close() - return get_json_result(retmsg="Database Connection Successful!") + return get_json_result(data="Database Connection Successful!") except Exception as e: return server_error_response(e)