From f86d8906e73b71fd2e384a35d2171f1da1f98fe0 Mon Sep 17 00:00:00 2001 From: so95 Date: Thu, 9 Jan 2025 10:55:18 +0700 Subject: [PATCH] Fixed code error when mssql returns multiple columns (#4420) Fixed code error when mssql returns multiple columns --- agent/component/exesql.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/agent/component/exesql.py b/agent/component/exesql.py index dd210ad36..265256124 100644 --- a/agent/component/exesql.py +++ b/agent/component/exesql.py @@ -112,8 +112,11 @@ class ExeSQL(Generate, ABC): 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: @@ -143,4 +146,4 @@ class ExeSQL(Generate, ABC): return None def debug(self, **kwargs): - return self._run([], **kwargs) \ No newline at end of file + return self._run([], **kwargs)