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
|