H commited on
Commit
add4a89
·
1 Parent(s): 842fb8a

Fix exeSQL component output (#2141)

Browse files

### What problem does this PR solve?

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. agent/component/exesql.py +7 -4
agent/component/exesql.py CHANGED
@@ -54,7 +54,7 @@ class ExeSQL(ComponentBase, ABC):
54
  setattr(self, "_loop", 0)
55
  if self._loop >= self._param.loop:
56
  self._loop = 0
57
- raise Exception("Maximum loop time exceeds. Can't query the correct data via sql statement.")
58
  self._loop += 1
59
 
60
  ans = self.get_input()
@@ -63,7 +63,7 @@ class ExeSQL(ComponentBase, ABC):
63
  ans = re.sub(r';.*?SELECT ', '; SELECT ', ans, flags=re.IGNORECASE)
64
  ans = re.sub(r';[^;]*$', r';', ans)
65
  if not ans:
66
- return ExeSQL.be_output("SQL statement not found!")
67
 
68
  if self._param.db_type in ["mysql", "mariadb"]:
69
  db = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
@@ -75,13 +75,16 @@ class ExeSQL(ComponentBase, ABC):
75
  try:
76
  db.connect()
77
  except Exception as e:
78
- return ExeSQL.be_output("**Error**: \nDatabase Connection Failed! \n" + str(e))
79
  sql_res = []
80
  for single_sql in re.split(r';', ans.replace(r"\n", " ")):
81
  if not single_sql:
82
  continue
83
  try:
84
  query = db.execute_sql(single_sql)
 
 
 
85
  single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
86
  single_res.columns = [i[0] for i in query.description]
87
  sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
@@ -91,6 +94,6 @@ class ExeSQL(ComponentBase, ABC):
91
  db.close()
92
 
93
  if not sql_res:
94
- return ExeSQL.be_output("No record in the database!")
95
 
96
  return pd.DataFrame(sql_res)
 
54
  setattr(self, "_loop", 0)
55
  if self._loop >= self._param.loop:
56
  self._loop = 0
57
+ raise Exception("Maximum loop time exceeds. Can't query the correct data via SQL statement.")
58
  self._loop += 1
59
 
60
  ans = self.get_input()
 
63
  ans = re.sub(r';.*?SELECT ', '; SELECT ', ans, flags=re.IGNORECASE)
64
  ans = re.sub(r';[^;]*$', r';', ans)
65
  if not ans:
66
+ raise Exception("SQL statement not found!")
67
 
68
  if self._param.db_type in ["mysql", "mariadb"]:
69
  db = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
 
75
  try:
76
  db.connect()
77
  except Exception as e:
78
+ raise Exception("Database Connection Failed! \n" + str(e))
79
  sql_res = []
80
  for single_sql in re.split(r';', ans.replace(r"\n", " ")):
81
  if not single_sql:
82
  continue
83
  try:
84
  query = db.execute_sql(single_sql)
85
+ if query.rowcount == 0:
86
+ sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n No record in the database!"})
87
+ continue
88
  single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
89
  single_res.columns = [i[0] for i in query.description]
90
  sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
 
94
  db.close()
95
 
96
  if not sql_res:
97
+ return ExeSQL.be_output("")
98
 
99
  return pd.DataFrame(sql_res)