mirror of
https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git
synced 2025-08-12 15:39:05 +08:00
Refactor component exesql (#2048)
### What problem does this PR solve? ### Type of change - [x] Refactoring --------- Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
This commit is contained in:
parent
f2c4d53c58
commit
404cdc0b6d
@ -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)
|
||||
|
Loading…
x
Reference in New Issue
Block a user