From f843dd05e5bc1e8ad4f65e1d98648f14098ee581 Mon Sep 17 00:00:00 2001
From: H <43509927+guoyuhao2330@users.noreply.github.com>
Date: Wed, 28 Aug 2024 16:39:10 +0800
Subject: [PATCH] Fix exeSQL component output (#2141)

### What problem does this PR solve?

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 agent/component/exesql.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/agent/component/exesql.py b/agent/component/exesql.py
index 39ad4d02c..2afb55cab 100644
--- a/agent/component/exesql.py
+++ b/agent/component/exesql.py
@@ -54,7 +54,7 @@ class ExeSQL(ComponentBase, ABC):
             setattr(self, "_loop", 0)
         if self._loop >= self._param.loop:
             self._loop = 0
-            raise Exception("Maximum loop time exceeds. Can't query the correct data via sql statement.")
+            raise Exception("Maximum loop time exceeds. Can't query the correct data via SQL statement.")
         self._loop += 1
 
         ans = self.get_input()
@@ -63,7 +63,7 @@ class ExeSQL(ComponentBase, ABC):
         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!")
+            raise Exception("SQL statement not found!")
 
         if self._param.db_type in ["mysql", "mariadb"]:
             db = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
@@ -75,13 +75,16 @@ class ExeSQL(ComponentBase, ABC):
         try:
             db.connect()
         except Exception as e:
-            return ExeSQL.be_output("**Error**: \nDatabase Connection Failed! \n" + str(e))
+            raise Exception("Database Connection Failed! \n" + str(e))
         sql_res = []
         for single_sql in re.split(r';', ans.replace(r"\n", " ")):
             if not single_sql:
                 continue
             try:
                 query = db.execute_sql(single_sql)
+                if query.rowcount == 0:
+                    sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n No record in the database!"})
+                    continue
                 single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
                 single_res.columns = [i[0] for i in query.description]
                 sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
@@ -91,6 +94,6 @@ class ExeSQL(ComponentBase, ABC):
         db.close()
 
         if not sql_res:
-            return ExeSQL.be_output("No record in the database!")
+            return ExeSQL.be_output("")
 
         return pd.DataFrame(sql_res)