Update exesql.py

agent/component/exesql.py

@@ -16,7 +16,8 @@
 from abc import ABC
 import re
 import pandas as pd
-from peewee import MySQLDatabase, PostgresqlDatabase
+import pymysql
+import psycopg2
 from agent.component.base import ComponentBase, ComponentParamBase
 
 
@@ -66,28 +67,28 @@ def _run(self, history, **kwargs):
             raise Exception("SQL statement not found!")
 
         if self._param.db_type in ["mysql", "mariadb"]:
-            db = MySQLDatabase(self._param.database, user=self._param.username, host=self._param.host,
-                               port=self._param.port, password=self._param.password)
+            db = pymysql.connect(db=self._param.database, user=self._param.username, host=self._param.host,
+                                 port=self._param.port, password=self._param.password)
         elif self._param.db_type == 'postgresql':
-            db = PostgresqlDatabase(self._param.database, user=self._param.username, host=self._param.host,
-                                    port=self._param.port, password=self._param.password)
+            db = psycopg2.connect(dbname=self._param.database, user=self._param.username, host=self._param.host,
+                                  port=self._param.port, password=self._param.password)
 
         try:
-            db.connect()
+            cursor = db.cursor()
         except Exception as e:
             raise Exception("Database Connection Failed! \n" + str(e))
         sql_res = []
         for single_sql in re.split(r';', ans.replace(r"\n", " ")):
             if not single_sql:
                 continue
             try:
-                query = db.execute_sql(single_sql)
-                if query.rowcount == 0:
-                    sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n No record in the database!"})
+                cursor.execute(single_sql)
+                if cursor.rowcount == 0:
+                    sql_res.append({"content": "\nTotal: 0\n No record in the database!"})
                     continue
-                single_res = pd.DataFrame([i for i in query.fetchmany(size=self._param.top_n)])
-                single_res.columns = [i[0] for i in query.description]
-                sql_res.append({"content": "\nTotal: " + str(query.rowcount) + "\n" + single_res.to_markdown()})
+                single_res = pd.DataFrame([i for i in cursor.fetchmany(size=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()})
             except Exception as e:
                 sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql})
                 pass