mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-09-27 12:33:16 +00:00
add connector file
This commit is contained in:
51
pilot/connections/mysql.py
Normal file
51
pilot/connections/mysql.py
Normal file
@@ -0,0 +1,51 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import pymysql
|
||||
|
||||
class MySQLOperator:
|
||||
"""Connect MySQL Database fetch MetaData For LLM Prompt
|
||||
Args:
|
||||
|
||||
Usage:
|
||||
"""
|
||||
|
||||
default_db = ["information_schema", "performance_schema", "sys", "mysql"]
|
||||
def __init__(self, user, password, host="localhost", port=3306) -> None:
|
||||
|
||||
self.conn = pymysql.connect(
|
||||
host=host,
|
||||
user=user,
|
||||
passwd=password,
|
||||
charset="utf8mb4",
|
||||
cursorclass=pymysql.cursors.DictCursor
|
||||
)
|
||||
|
||||
def get_schema(self, schema_name):
|
||||
|
||||
with self.conn.cursor() as cursor:
|
||||
_sql = f"""
|
||||
select concat(table_name, "(" , group_concat(column_name), ")") as schema_info from information_schema.COLUMNS where table_schema="{schema_name}" group by TABLE_NAME;
|
||||
"""
|
||||
cursor.execute(_sql)
|
||||
results = cursor.fetchall()
|
||||
return results
|
||||
|
||||
def get_index(self, schema_name):
|
||||
pass
|
||||
|
||||
def get_db_list(self):
|
||||
with self.conn.cursor() as cursor:
|
||||
_sql = """
|
||||
show databases;
|
||||
"""
|
||||
cursor.execute(_sql)
|
||||
results = cursor.fetchall()
|
||||
|
||||
dbs = [d["Database"] for d in results if d["Database"] not in self.default_db]
|
||||
return dbs
|
||||
|
||||
def get_meta(self, schema_name):
|
||||
pass
|
||||
|
||||
|
Reference in New Issue
Block a user