Skip to content

Commit

Permalink
Fixed code error when mssql returns multiple columns (#4420)
Browse files Browse the repository at this point in the history
Fixed code error when mssql returns multiple columns
  • Loading branch information
isthaison authored Jan 9, 2025
1 parent bc681e2 commit f86d890
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions agent/component/exesql.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,11 @@ def _run(self, history, **kwargs):
if cursor.rowcount == 0:
sql_res.append({"content": "\nTotal: 0\n No record in the database!"})
break
single_res = pd.DataFrame([i for i in cursor.fetchmany(self._param.top_n)])
single_res.columns = [i[0] for i in cursor.description]
if self._param.db_type == 'mssql':
single_res = pd.DataFrame.from_records(cursor.fetchmany(self._param.top_n),columns = [desc[0] for desc in cursor.description])
else:
single_res = pd.DataFrame([i for i in cursor.fetchmany(self._param.top_n)])
single_res.columns = [i[0] for i in cursor.description]
sql_res.append({"content": "\nTotal: " + str(cursor.rowcount) + "\n" + single_res.to_markdown()})
break
except Exception as e:
Expand Down Expand Up @@ -143,4 +146,4 @@ def _regenerate_sql(self, failed_sql, error_message,**kwargs):
return None

def debug(self, **kwargs):
return self._run([], **kwargs)
return self._run([], **kwargs)

0 comments on commit f86d890

Please sign in to comment.