386 lines
13 KiB
Python
386 lines
13 KiB
Python
import os
|
|
import sys
|
|
import asyncio
|
|
from pathlib import Path
|
|
import aiomysql
|
|
import datetime
|
|
import inspect
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from config import settings
|
|
from database.models import Base
|
|
from sqlalchemy import inspect as sqlalchemy_inspect
|
|
from utils.logger import session_logger
|
|
|
|
async def get_mysql_connection():
|
|
"""建立MySQL数据库连接"""
|
|
try:
|
|
conn = await aiomysql.connect(
|
|
host=settings.db_host,
|
|
port=settings.db_port,
|
|
user=settings.db_user,
|
|
password=settings.db_password,
|
|
db=settings.db_name,
|
|
autocommit=False
|
|
)
|
|
return conn
|
|
except Exception as e:
|
|
print(f"数据库连接失败: {e}")
|
|
return None
|
|
|
|
async def get_all_tables(conn):
|
|
"""获取数据库中所有表"""
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("SHOW TABLES")
|
|
tables = await cursor.fetchall()
|
|
return [table[0] for table in tables]
|
|
|
|
async def get_model_tables():
|
|
"""获取模型定义的所有表"""
|
|
return [table.__tablename__ for table in Base.__subclasses__()]
|
|
|
|
async def get_table_columns(conn, table_name):
|
|
"""获取表的所有列"""
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
|
|
columns = await cursor.fetchall()
|
|
return [column[0] for column in columns]
|
|
|
|
async def get_model_columns(model_class):
|
|
"""获取模型定义的所有列"""
|
|
inspector = sqlalchemy_inspect(model_class)
|
|
return [column.key for column in inspector.mapper.column_attrs]
|
|
|
|
async def identify_unused_tables(conn):
|
|
"""识别未使用的表"""
|
|
db_tables = await get_all_tables(conn)
|
|
model_tables = await get_model_tables()
|
|
|
|
unused_tables = [table for table in db_tables if table not in model_tables]
|
|
return unused_tables
|
|
|
|
async def identify_unused_columns(conn):
|
|
"""识别未使用的列"""
|
|
results = {}
|
|
model_dict = {model.__tablename__: model for model in Base.__subclasses__()}
|
|
|
|
for table_name, model_class in model_dict.items():
|
|
try:
|
|
db_columns = await get_table_columns(conn, table_name)
|
|
model_columns = await get_model_columns(model_class)
|
|
|
|
unused_columns = [col for col in db_columns if col not in model_columns]
|
|
if unused_columns:
|
|
results[table_name] = unused_columns
|
|
except Exception as e:
|
|
print(f"检查表 {table_name} 列时出错: {e}")
|
|
|
|
return results
|
|
|
|
async def get_table_foreign_keys(conn, table_name):
|
|
"""获取表的外键约束"""
|
|
async with conn.cursor() as cursor:
|
|
query = """
|
|
SELECT
|
|
CONSTRAINT_NAME,
|
|
TABLE_NAME,
|
|
COLUMN_NAME,
|
|
REFERENCED_TABLE_NAME,
|
|
REFERENCED_COLUMN_NAME
|
|
FROM
|
|
INFORMATION_SCHEMA.KEY_COLUMN_USAGE
|
|
WHERE
|
|
REFERENCED_TABLE_NAME IS NOT NULL
|
|
AND TABLE_SCHEMA = %s
|
|
AND TABLE_NAME = %s
|
|
"""
|
|
await cursor.execute(query, (settings.db_name, table_name))
|
|
foreign_keys = await cursor.fetchall()
|
|
return foreign_keys
|
|
|
|
async def get_referenced_tables(conn, table_name):
|
|
"""获取引用指定表的所有表"""
|
|
async with conn.cursor() as cursor:
|
|
query = """
|
|
SELECT
|
|
TABLE_NAME,
|
|
CONSTRAINT_NAME
|
|
FROM
|
|
INFORMATION_SCHEMA.KEY_COLUMN_USAGE
|
|
WHERE
|
|
REFERENCED_TABLE_NAME = %s
|
|
AND TABLE_SCHEMA = %s
|
|
"""
|
|
await cursor.execute(query, (table_name, settings.db_name))
|
|
references = await cursor.fetchall()
|
|
|
|
if references:
|
|
print(f"表 {table_name} 被以下表引用:")
|
|
for ref in references:
|
|
print(f" - {ref[0]} (约束: {ref[1]})")
|
|
|
|
return references
|
|
|
|
async def disable_foreign_key_checks(conn):
|
|
"""临时禁用外键约束检查"""
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
|
|
print("已临时禁用外键约束检查")
|
|
|
|
async def enable_foreign_key_checks(conn):
|
|
"""重新启用外键约束检查"""
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
|
|
print("已重新启用外键约束检查")
|
|
|
|
async def sort_tables_for_deletion(conn, tables):
|
|
"""对表进行排序，以确保先删除引用者，再删除被引用者"""
|
|
dependency_graph = {}
|
|
for table in tables:
|
|
references = await get_referenced_tables(conn, table)
|
|
dependency_graph[table] = [ref[0] for ref in references if ref[0] in tables]
|
|
|
|
result = []
|
|
visited = set()
|
|
temp_mark = set()
|
|
|
|
def visit(node):
|
|
if node in temp_mark:
|
|
return
|
|
if node not in visited:
|
|
temp_mark.add(node)
|
|
for dependent in dependency_graph.get(node, []):
|
|
visit(dependent)
|
|
temp_mark.remove(node)
|
|
visited.add(node)
|
|
result.append(node)
|
|
|
|
for table in tables:
|
|
if table not in visited:
|
|
visit(table)
|
|
|
|
return result[::-1]
|
|
|
|
async def drop_tables(conn, tables, dry_run=True):
|
|
"""删除未使用的表"""
|
|
if not tables:
|
|
return True
|
|
|
|
if not dry_run:
|
|
try:
|
|
print("禁用外键约束检查以确保安全删除所有表...")
|
|
await disable_foreign_key_checks(conn)
|
|
|
|
for table in tables:
|
|
try:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(f"DROP TABLE IF EXISTS `{table}`")
|
|
print(f"已删除表: {table}")
|
|
except Exception as e:
|
|
print(f"删除表 {table} 时出错: {e}")
|
|
await conn.rollback()
|
|
await enable_foreign_key_checks(conn)
|
|
return False
|
|
|
|
await conn.commit()
|
|
await enable_foreign_key_checks(conn)
|
|
return True
|
|
except Exception as e:
|
|
print(f"删除表时出错: {e}")
|
|
try:
|
|
await enable_foreign_key_checks(conn)
|
|
await conn.rollback()
|
|
except:
|
|
pass
|
|
return False
|
|
else:
|
|
for table in tables:
|
|
print(f"将删除表: {table}")
|
|
return True
|
|
|
|
async def drop_columns(conn, column_dict, dry_run=True):
|
|
"""删除未使用的列"""
|
|
for table, columns in column_dict.items():
|
|
try:
|
|
for column in columns:
|
|
if dry_run:
|
|
print(f"将从表 {table} 删除列: {column}")
|
|
else:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(f"ALTER TABLE `{table}` DROP COLUMN `{column}`")
|
|
print(f"已从表 {table} 删除列: {column}")
|
|
except Exception as e:
|
|
print(f"从表 {table} 删除列时出错: {e}")
|
|
if not dry_run:
|
|
await conn.rollback()
|
|
return False
|
|
|
|
if not dry_run:
|
|
await conn.commit()
|
|
return True
|
|
|
|
async def cleanup_database(dry_run=True):
|
|
"""清理数据库中未使用的表和列"""
|
|
conn = await get_mysql_connection()
|
|
if not conn:
|
|
return False
|
|
|
|
try:
|
|
unused_tables = await identify_unused_tables(conn)
|
|
unused_columns = await identify_unused_columns(conn)
|
|
|
|
print("\n=== 数据库清理报告 ===")
|
|
|
|
if unused_tables:
|
|
print(f"\n发现 {len(unused_tables)} 个未使用的表:")
|
|
for table in unused_tables:
|
|
print(f" - {table}")
|
|
|
|
if not dry_run:
|
|
success = await drop_tables(conn, unused_tables, dry_run=False)
|
|
if not success:
|
|
print("删除表失败，已回滚操作")
|
|
return False
|
|
else:
|
|
print("\n未发现任何未使用的表")
|
|
|
|
if unused_columns:
|
|
total_columns = sum(len(cols) for cols in unused_columns.values())
|
|
print(f"\n发现 {total_columns} 个未使用的列:")
|
|
for table, columns in unused_columns.items():
|
|
for column in columns:
|
|
print(f" - {table}.{column}")
|
|
|
|
if not dry_run:
|
|
success = await drop_columns(conn, unused_columns, dry_run=False)
|
|
if not success:
|
|
print("删除列失败，已回滚操作")
|
|
return False
|
|
else:
|
|
print("\n未发现任何未使用的列")
|
|
|
|
if dry_run:
|
|
print("\n这是一次模拟运行，未执行任何实际更改。要执行更改，请使用 --execute 选项。")
|
|
else:
|
|
print("\n数据库清理成功完成！")
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f"清理数据库时发生错误: {e}")
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
async def remove_orphaned_records(dry_run=True):
|
|
"""删除孤立记录（有外键但对应的父记录不存在）"""
|
|
conn = await get_mysql_connection()
|
|
if not conn:
|
|
return False
|
|
|
|
try:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("""
|
|
SELECT COUNT(*) FROM conversation_messages
|
|
WHERE conversation_id NOT IN (SELECT id FROM conversation_histories)
|
|
""")
|
|
count = (await cursor.fetchone())[0]
|
|
|
|
if count > 0:
|
|
print(f"\n发现 {count} 条孤立的会话消息记录")
|
|
if not dry_run:
|
|
await cursor.execute("""
|
|
DELETE FROM conversation_messages
|
|
WHERE conversation_id NOT IN (SELECT id FROM conversation_histories)
|
|
""")
|
|
print(f"已删除 {count} 条孤立的会话消息记录")
|
|
else:
|
|
print("\n未发现孤立的会话消息记录")
|
|
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute("""
|
|
SELECT COUNT(*) FROM role_languages
|
|
WHERE role_id NOT IN (SELECT id FROM roles)
|
|
""")
|
|
count = (await cursor.fetchone())[0]
|
|
|
|
if count > 0:
|
|
print(f"\n发现 {count} 条孤立的角色语言记录")
|
|
if not dry_run:
|
|
await cursor.execute("""
|
|
DELETE FROM role_languages
|
|
WHERE role_id NOT IN (SELECT id FROM roles)
|
|
""")
|
|
print(f"已删除 {count} 条孤立的角色语言记录")
|
|
else:
|
|
print("\n未发现孤立的角色语言记录")
|
|
|
|
if not dry_run:
|
|
await conn.commit()
|
|
|
|
return True
|
|
except Exception as e:
|
|
print(f"清理孤立记录时发生错误: {e}")
|
|
if not dry_run:
|
|
await conn.rollback()
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
async def drop_specific_tables(conn, tables_to_drop, dry_run=True):
|
|
"""删除指定的表，无论它们是否被使用"""
|
|
if dry_run:
|
|
print(f"\n将删除以下指定表:")
|
|
for table in tables_to_drop:
|
|
print(f" - {table}")
|
|
return True
|
|
|
|
try:
|
|
await disable_foreign_key_checks(conn)
|
|
|
|
for table in tables_to_drop:
|
|
try:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(f"DROP TABLE IF EXISTS `{table}`")
|
|
print(f"已删除表: {table}")
|
|
except Exception as e:
|
|
print(f"删除表 {table} 时出错: {e}")
|
|
|
|
await conn.commit()
|
|
print(f"指定表已成功删除")
|
|
return True
|
|
except Exception as e:
|
|
await conn.rollback()
|
|
print(f"删除指定表时出错: {e}")
|
|
return False
|
|
finally:
|
|
await enable_foreign_key_checks(conn)
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="数据库清理工具")
|
|
parser.add_argument('--execute', action='store_true', help='执行实际更改（不使用此选项则只进行模拟运行）')
|
|
parser.add_argument('--drop-tables', type=str, help='指定要删除的表，用逗号分隔')
|
|
args = parser.parse_args()
|
|
|
|
dry_run = not args.execute
|
|
|
|
async def main():
|
|
if args.drop_tables:
|
|
conn = await get_mysql_connection()
|
|
if not conn:
|
|
return False
|
|
|
|
try:
|
|
tables_to_drop = [t.strip() for t in args.drop_tables.split(',')]
|
|
await drop_specific_tables(conn, tables_to_drop, dry_run)
|
|
finally:
|
|
conn.close()
|
|
else:
|
|
result = await cleanup_database(dry_run)
|
|
if result:
|
|
await remove_orphaned_records(dry_run)
|
|
|
|
asyncio.run(main())
|