from .base import MetadataProvider from typing import List, Dict, Any from sqlalchemy import create_engine, text, inspect class MySQLMetadataProvider(MetadataProvider): def __init__(self, config: Dict): super().__init__(config) url = f"mysql+pymysql://{config['username']}:{config.get('password','')}@{config['host']}:{config['port']}/{config.get('database','')}?charset=utf8mb4" self.engine = create_engine(url, echo=True ) # 设为 True 可打印 SQL 日志 self.inspector = inspect(self.engine) def _get_connection(self): return self.engine.connect() # return pymysql.connect( # host=self.config["host"], # port=self.config["port"], # user=self.config["username"], # password=self.config.get("password", ""), # database=self.config.get("database", ""), # 可为空 # charset='utf8mb4' # ) def get_databases(self) -> List[str]: # with self._get_connection() as conn: # # cursor = conn.cursor() # # cursor.execute("SHOW DATABASES") # # return [row[0] for row in cursor.fetchall() if not row[0].startswith('information_schema')] # result = conn.execute(text("SHOW DATABASES")) # return [row[0] for row in result if not row[0].startswith('information_schema')] return [row for row in self.inspector.get_schema_names() if not row.startswith('information_schema')] def get_tables(self, database: str) -> List[str]: return self.inspector.get_table_names(schema=database) def get_columns(self, table: str, database: str) -> List[Dict[str, Any]]: return self.inspector.get_columns(table_name=table, schema=database) # with self._get_connection() as conn: # cursor = conn.cursor() # # 查询 information_schema 获取字段详情 # query = """ # SELECT COLUMN_NAME, DATA_TYPE, COLUMN_COMMENT # FROM information_schema.COLUMNS # WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s # ORDER BY ORDINAL_POSITION # """ # cursor.execute(query, (database, table)) # return [ # { # "name": row[0], # "type": row[1], # "comment": row[2] or "" # } # for row in cursor.fetchall() # ] def has_table(self, table_name:str, schema:str) -> bool: return self.inspector.has_table(table_name=table_name, schema=schema) def get_table_comment(self, table_name:str, schema:str) -> str: return self.inspector.get_table_comment(table_name=table_name, schema=schema)['text'] def get_pk_constraint(self, table_name:str, schema:str) -> Dict: return self.inspector.get_pk_constraint(table_name=table_name, schema=schema) def get_unique_constraints(self, table_name:str, schema:str) -> Dict: return self.inspector.get_unique_constraints(table_name=table_name, schema=schema) def get_foreign_keys(self, table_name:str, schema:str) -> list[Dict]: return self.inspector.get_foreign_keys(table_name=table_name, schema=schema) def fectch_distinct_values(self, column_name:str, table_name:str, schema:str, max_num:int = 5) -> list[str]: with self._get_connection() as conn: result = conn.execute(text(f"select distinct `{column_name}` from `{schema}`.`{table_name}` where `{column_name}` is not null limit {max_num}")) return [row[0] for row in result]