Browse Source

Fix component exesql (#2754)

### What problem does this PR solve?

#2700 

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
tags/v0.13.0
lidp 1 year ago
parent
commit
8f815a6c1e
No account linked to committer's email address
1 changed files with 13 additions and 12 deletions
  1. 13
    12
      agent/component/exesql.py

+ 13
- 12
agent/component/exesql.py View File

from abc import ABC from abc import ABC
import re import re
import pandas as pd import pandas as pd
from peewee import MySQLDatabase, PostgresqlDatabase
import pymysql
import psycopg2
from agent.component.base import ComponentBase, ComponentParamBase from agent.component.base import ComponentBase, ComponentParamBase




raise Exception("SQL statement not found!") raise Exception("SQL statement not found!")


if self._param.db_type in ["mysql", "mariadb"]: 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': 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: try:
db.connect()
cursor = db.cursor()
except Exception as e: except Exception as e:
raise Exception("Database Connection Failed! \n" + str(e)) raise Exception("Database Connection Failed! \n" + str(e))
sql_res = [] sql_res = []
if not single_sql: if not single_sql:
continue continue
try: 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 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: except Exception as e:
sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql}) sql_res.append({"content": "**Error**:" + str(e) + "\nError SQL Statement:" + single_sql})
pass pass

Loading…
Cancel
Save