H Kevin Hu commited on
Commit
4eefe86
·
1 Parent(s): 5fcb7d4

Refactor component exesql (#2048)

Browse files

### What problem does this PR solve?

### Type of change

- [x] Refactoring

---------

Co-authored-by: Kevin Hu <[email protected]>

Files changed (1) hide show
  1. agent/component/exesql.py +15 -7
agent/component/exesql.py CHANGED
@@ -74,15 +74,23 @@ class ExeSQL(ComponentBase, ABC):
74
 
75
  try:
76
  db.connect()
77
- query = db.execute_sql(ans)
78
- sql_res = [{"content": rec + "\n"} for rec in [str(i) for i in query.fetchall()]]
79
- db.close()
80
  except Exception as e:
81
- return ExeSQL.be_output("**Error**:" + str(e) + "\nError SQL Statement:" + ans)
 
 
 
 
 
 
 
 
 
 
 
 
 
82
 
83
  if not sql_res:
84
  return ExeSQL.be_output("No record in the database!")
85
 
86
- sql_res.insert(0, {"content": "Number of records retrieved from the database is " + str(len(sql_res)) + "\n"})
87
- df = pd.DataFrame(sql_res[0:self._param.top_n + 1])
88
- return ExeSQL.be_output(df.to_markdown())
 
74
 
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):
81
+ if not single_sql:
82
+ continue
83
+ try:
84
+ query = db.execute_sql(single_sql)
85
+ sql_res.append(
86
+ {"content": "\n##Total: " + str(query.rowcount) + "\n" + pd.DataFrame(
87
+ [i for i in query.fetchmany(size=self._param.top_n)]).to_markdown()})
88
+ except Exception as e:
89
+ sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql})
90
+ pass
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)