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)

Files changed (1) hide show
  1. 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
- from peewee import MySQLDatabase, PostgresqlDatabase
 
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 = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
70
- port=self._param.port, password=self._param.password)
71
  elif self._param.db_type == 'postgresql':
72
- db = PostgresqlDatabase(self._param.database, user=self._param.username, host=self._param.host,
73
- port=self._param.port, password=self._param.password)
74
 
75
  try:
76
- db.connect()
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
- query = db.execute_sql(single_sql)
85
- if query.rowcount == 0:
86
- sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n No record in the database!"})
87
  continue
88
- single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
89
- single_res.columns = [i[0] for i in query.description]
90
- sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
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