isthaison commited on
Commit
96bcfec
·
1 Parent(s): 1c17452

Fixed code error when mssql returns multiple columns (#4420)

Browse files
Files changed (1) hide show
  1. agent/component/exesql.py +6 -3
agent/component/exesql.py CHANGED
@@ -112,8 +112,11 @@ class ExeSQL(Generate, ABC):
112
  if cursor.rowcount == 0:
113
  sql_res.append({"content": "\nTotal: 0\n No record in the database!"})
114
  break
115
- single_res = pd.DataFrame([i for i in cursor.fetchmany(self._param.top_n)])
116
- single_res.columns = [i[0] for i in cursor.description]
 
 
 
117
  sql_res.append({"content": "\nTotal: " + str(cursor.rowcount) + "\n" + single_res.to_markdown()})
118
  break
119
  except Exception as e:
@@ -143,4 +146,4 @@ class ExeSQL(Generate, ABC):
143
  return None
144
 
145
  def debug(self, **kwargs):
146
- return self._run([], **kwargs)
 
112
  if cursor.rowcount == 0:
113
  sql_res.append({"content": "\nTotal: 0\n No record in the database!"})
114
  break
115
+ if self._param.db_type == 'mssql':
116
+ single_res = pd.DataFrame.from_records(cursor.fetchmany(self._param.top_n),columns = [desc[0] for desc in cursor.description])
117
+ else:
118
+ single_res = pd.DataFrame([i for i in cursor.fetchmany(self._param.top_n)])
119
+ single_res.columns = [i[0] for i in cursor.description]
120
  sql_res.append({"content": "\nTotal: " + str(cursor.rowcount) + "\n" + single_res.to_markdown()})
121
  break
122
  except Exception as e:
 
146
  return None
147
 
148
  def debug(self, **kwargs):
149
+ return self._run([], **kwargs)