lidp
commited on
Commit
·
db99ead
1
Parent(s):
d607735
Fix component exesql (#2754)
Browse files### What problem does this PR solve?
#2700
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- agent/component/exesql.py +13 -12
agent/component/exesql.py
CHANGED
@@ -16,7 +16,8 @@
|
|
16 |
from abc import ABC
|
17 |
import re
|
18 |
import pandas as pd
|
19 |
-
|
|
|
20 |
from agent.component.base import ComponentBase, ComponentParamBase
|
21 |
|
22 |
|
@@ -66,14 +67,14 @@ class ExeSQL(ComponentBase, ABC):
|
|
66 |
raise Exception("SQL statement not found!")
|
67 |
|
68 |
if self._param.db_type in ["mysql", "mariadb"]:
|
69 |
-
db =
|
70 |
-
|
71 |
elif self._param.db_type == 'postgresql':
|
72 |
-
db =
|
73 |
-
|
74 |
|
75 |
try:
|
76 |
-
db.
|
77 |
except Exception as e:
|
78 |
raise Exception("Database Connection Failed! \n" + str(e))
|
79 |
sql_res = []
|
@@ -81,13 +82,13 @@ class ExeSQL(ComponentBase, ABC):
|
|
81 |
if not single_sql:
|
82 |
continue
|
83 |
try:
|
84 |
-
|
85 |
-
if
|
86 |
-
sql_res.append({"content": "\nTotal:
|
87 |
continue
|
88 |
-
single_res = pd.DataFrame([i for i in
|
89 |
-
single_res.columns = [i[0] for i in
|
90 |
-
sql_res.append({"content": "\nTotal: " + str(
|
91 |
except Exception as e:
|
92 |
sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql})
|
93 |
pass
|
|
|
16 |
from abc import ABC
|
17 |
import re
|
18 |
import pandas as pd
|
19 |
+
import pymysql
|
20 |
+
import psycopg2
|
21 |
from agent.component.base import ComponentBase, ComponentParamBase
|
22 |
|
23 |
|
|
|
67 |
raise Exception("SQL statement not found!")
|
68 |
|
69 |
if self._param.db_type in ["mysql", "mariadb"]:
|
70 |
+
db = pymysql.connect(db=self._param.database, user=self._param.username, host=self._param.host,
|
71 |
+
port=self._param.port, password=self._param.password)
|
72 |
elif self._param.db_type == 'postgresql':
|
73 |
+
db = psycopg2.connect(dbname=self._param.database, user=self._param.username, host=self._param.host,
|
74 |
+
port=self._param.port, password=self._param.password)
|
75 |
|
76 |
try:
|
77 |
+
cursor = db.cursor()
|
78 |
except Exception as e:
|
79 |
raise Exception("Database Connection Failed! \n" + str(e))
|
80 |
sql_res = []
|
|
|
82 |
if not single_sql:
|
83 |
continue
|
84 |
try:
|
85 |
+
cursor.execute(single_sql)
|
86 |
+
if cursor.rowcount == 0:
|
87 |
+
sql_res.append({"content": "\nTotal: 0\n No record in the database!"})
|
88 |
continue
|
89 |
+
single_res = pd.DataFrame([i for i in cursor.fetchmany(size=self._param.top_n)])
|
90 |
+
single_res.columns = [i[0] for i in cursor.description]
|
91 |
+
sql_res.append({"content": "\nTotal: " + str(cursor.rowcount) + "\n" + single_res.to_markdown()})
|
92 |
except Exception as e:
|
93 |
sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql})
|
94 |
pass
|