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]>
- 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**:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
|
83 |
if not sql_res:
|
84 |
return ExeSQL.be_output("No record in the database!")
|
85 |
|
86 |
-
|
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)
|
|
|
|