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 ''}")