add connector file

This commit is contained in:
csunny
2023-05-08 00:44:02 +08:00
parent d746086694
commit 3ce31d5ea0
9 changed files with 59 additions and 3 deletions

View 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