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)
|
|
|
|
|
|