diff --git a/agent/component/exesql.py b/agent/component/exesql.py index a5ce09e74..3aea7a8ae 100644 --- a/agent/component/exesql.py +++ b/agent/component/exesql.py @@ -74,15 +74,23 @@ class ExeSQL(ComponentBase, ABC): try: db.connect() - query = db.execute_sql(ans) - sql_res = [{"content": rec + "\n"} for rec in [str(i) for i in query.fetchall()]] - db.close() except Exception as e: - return ExeSQL.be_output("**Error**:" + str(e) + "\nError SQL Statement:" + ans) + return ExeSQL.be_output("**Error**: \nDatabase Connection Failed! \n" + str(e)) + sql_res = [] + for single_sql in re.split(r';', ans): + if not single_sql: + continue + try: + query = db.execute_sql(single_sql) + sql_res.append( + {"content": "\n##Total: " + str(query.rowcount) + "\n" + pd.DataFrame( + [i for i in query.fetchmany(size=self._param.top_n)]).to_markdown()}) + except Exception as e: + sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql}) + pass + db.close() if not sql_res: return ExeSQL.be_output("No record in the database!") - sql_res.insert(0, {"content": "Number of records retrieved from the database is " + str(len(sql_res)) + "\n"}) - df = pd.DataFrame(sql_res[0:self._param.top_n + 1]) - return ExeSQL.be_output(df.to_markdown()) + return pd.DataFrame(sql_res)