import pymysql
from pymysql.cursors import DictCursor
from typing import Optional, Dict, List, Anyclass DBSchemaExporter:"""MySQL数据库表结构导出工具，支持提取表和字段注释使用示例:>>> exporter = DBSchemaExporter("localhost", "user", "password", "dbname")>>> schema = exporter.export(include_comments=True)>>> print(schema)>>> exporter.save_to_file("schema.txt")"""def __init__(self, host: str, user: str, password: str, database: str, port: int = 3306, charset: str = 'utf8mb4'):"""初始化数据库连接配置"""self.config = {'host': host,'user': user,'password': password,'database': database,'port': port,'charset': charset,'cursorclass': DictCursor}self.schema = Nonedef export(self, include_comments: bool = True, include_tables: Optional[List[str]] = None, exclude_tables: Optional[List[str]] = None) -> Optional[str]:"""导出数据库表结构，支持包含注释Args:include_comments: 是否包含表和字段注释include_tables: 只包含指定的表exclude_tables: 排除指定的表Returns:包含所有表结构的字符串，失败时返回None"""try:# 获取表和字段注释信息table_comments = {}column_comments = {}if include_comments:table_comments, column_comments = self._fetch_comments()with pymysql.connect(**self.config) as connection:with connection.cursor() as cursor:# 获取所有表名cursor.execute("SHOW TABLES")all_tables = [row[f"Tables_in_{self.config['database']}"] for row in cursor.fetchall()]# 过滤表if include_tables:tables = [t for t in all_tables if t in include_tables]else:tables = all_tablesif exclude_tables:tables = [t for t in tables if t not in exclude_tables]# 获取每个表的结构table_schemas = []for table in tables:# 获取表的创建语句cursor.execute(f"SHOW CREATE TABLE `{table}`")create_table = cursor.fetchone()["Create Table"]# 提取表结构部分table_structure = create_table.split(f"CREATE TABLE `{table}`", 1)[1].strip()# 添加表注释（如果有）if include_comments and table in table_comments:comment_line = f"-- 表注释: {table_comments[table]}"table_schemas.append(f"{comment_line}\n{table} (\n{table_structure}\n)\n")else:table_schemas.append(f"{table} (\n{table_structure}\n)\n")self.schema = "\n\n".join(table_schemas)return self.schemaexcept Exception as e:print(f"导出失败: {e}")return Nonedef _fetch_comments(self) -> tuple:"""获取所有表和字段的注释信息"""table_comments = {}column_comments = {}with pymysql.connect(**self.config) as connection:with connection.cursor() as cursor:# 获取表注释cursor.execute("""SELECT TABLE_NAME, TABLE_COMMENT FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = %s""", (self.config['database'],))for row in cursor.fetchall():table_comments[row['TABLE_NAME']] = row['TABLE_COMMENT']# 获取字段注释cursor.execute("""SELECT TABLE_NAME, COLUMN_NAME, COLUMN_COMMENT FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = %s""", (self.config['database'],))for row in cursor.fetchall():table = row['TABLE_NAME']column = row['COLUMN_NAME']comment = row['COLUMN_COMMENT']if table not in column_comments:column_comments[table] = {}column_comments[table][column] = commentreturn table_comments, column_commentsdef save_to_file(self, file_path: str, overwrite: bool = False) -> bool:"""将导出的表结构保存到文件Args:file_path: 文件路径overwrite: 是否覆盖已存在的文件Returns:保存成功返回True，失败返回False"""if self.schema is None:print("没有导出的表结构，请先调用export()方法")return Falsetry:# 检查文件是否存在import osif os.path.exists(file_path) and not overwrite:print(f"文件已存在: {file_path}，设置overwrite=True以覆盖")return Falsewith open(file_path, 'w', encoding='utf-8') as f:f.write(self.schema)return Trueexcept Exception as e:print(f"保存失败: {e}")return Falsedef get_table_details(self) -> Optional[Dict[str, List[Dict[str, Any]]]]:"""获取每个表的详细列信息，包括注释Returns:包含表和列信息的字典，格式为:{'table1': [{'Field': 'id', 'Type': 'int', 'Comment': '主键'}, ...],'table2': [...]}"""try:table_details = {}with pymysql.connect(**self.config) as connection:with connection.cursor() as cursor:# 获取所有表名cursor.execute("SHOW TABLES")tables = [row[f"Tables_in_{self.config['database']}"] for row in cursor.fetchall()]for table in tables:cursor.execute(f"SHOW FULL COLUMNS FROM `{table}`")columns = cursor.fetchall()table_details[table] = columnsreturn table_detailsexcept Exception as e:print(f"获取表详细信息失败: {e}")return None# 使用示例
if __name__ == "__main__":# 配置数据库连接exporter = DBSchemaExporter(host="localhost",user="your_username",password="your_password",database="your_database")# 示例1: 导出包含注释的完整表结构schema_with_comments = exporter.export(include_comments=True)if schema_with_comments:print("带注释的完整表结构:")print(schema_with_comments)exporter.save_to_file("full_schema_with_comments.txt", overwrite=True)# 示例2: 导出特定表的结构（不带注释）specific_schema = exporter.export(include_comments=False,include_tables=["users", "orders"])if specific_schema:print("\n特定表的结构（不带注释）:")print(specific_schema)exporter.save_to_file("specific_schema.txt", overwrite=True)# 示例3: 获取详细的列信息（包括注释）table_details = exporter.get_table_details()if table_details:print("\n表列详细信息:")for table, columns in table_details.items():print(f"\n{table}:")for column in columns[:3]:  # 只显示每个表的前3列comment = column.get('Comment', '')print(f"  - {column['Field']} ({column['Type']}){' - ' + comment if comment else ''}")