Merge branch 'dev' of https://github.com/MoFox-Studio/MoFox_Bot into dev

src/common/database/db_migration.py

@@ -1,6 +1,8 @@
 # mmc/src/common/database/db_migration.py
 
-from sqlalchemy import inspect, text
+from sqlalchemy import inspect
+from sqlalchemy.schema import AddColumn, CreateIndex
+
 from src.common.database.sqlalchemy_models import Base, get_engine
 from src.common.logger import get_logger
 
@@ -9,79 +11,106 @@ logger = get_logger("db_migration")
 
 async def check_and_migrate_database():
     """
-    异步检查数据库结构并自动迁移（添加缺失的表和列）。
+    异步检查数据库结构并自动迁移。
+    - 自动创建不存在的表。
+    - 自动为现有表添加缺失的列。
+    - 自动为现有表创建缺失的索引。
     """
     logger.info("正在检查数据库结构并执行自动迁移...")
     engine = await get_engine()
-    
-    # 使用异步引擎获取inspector
+
     async with engine.connect() as connection:
         # 在同步上下文中运行inspector操作
-        inspector = await connection.run_sync(lambda sync_conn: inspect(sync_conn))
-        
-        # 1. 获取数据库中所有已存在的表名
-        db_table_names = await connection.run_sync(lambda sync_conn: set(inspect(sync_conn).get_table_names()))
+        def get_inspector(sync_conn):
+            return inspect(sync_conn)
 
-        # 2. 遍历所有在代码中定义的模型
+        inspector = await connection.run_sync(get_inspector)
+
+        # 在同步lambda中传递inspector
+        db_table_names = await connection.run_sync(lambda conn: set(inspector.get_table_names(conn)))
+
+        # 1. 首先处理表的创建
+        tables_to_create = []
         for table_name, table in Base.metadata.tables.items():
-            logger.debug(f"正在检查表: {table_name}")
-
-            # 3. 如果表不存在，则创建它
             if table_name not in db_table_names:
-                logger.info(f"表 '{table_name}' 不存在，正在创建...")
-                try:
-                    await connection.run_sync(lambda sync_conn: table.create(sync_conn))
-                    logger.info(f"表 '{table_name}' 创建成功。")
-                except Exception as e:
-                    logger.error(f"创建表 '{table_name}' 失败: {e}")
+                tables_to_create.append(table)
+
+        if tables_to_create:
+            logger.info(f"发现 {len(tables_to_create)} 个不存在的表，正在创建...")
+            try:
+                # 一次性创建所有缺失的表
+                await connection.run_sync(
+                    lambda sync_conn: Base.metadata.create_all(sync_conn, tables=tables_to_create)
+                )
+                for table in tables_to_create:
+                    logger.info(f"表 '{table.name}' 创建成功。")
+                    db_table_names.add(table.name)  # 将新创建的表添加到集合中
+            except Exception as e:
+                logger.error(f"创建表时失败: {e}", exc_info=True)
+
+        # 2. 然后处理现有表的列和索引的添加
+        for table_name, table in Base.metadata.tables.items():
+            if table_name not in db_table_names:
+                logger.warning(f"跳过检查表 '{table_name}'，因为它在创建步骤中可能已失败。")
                 continue
 
-            # 4. 如果表已存在，则检查并添加缺失的列
-            db_columns = await connection.run_sync(
-                lambda sync_conn: {col["name"] for col in inspect(sync_conn).get_columns(table_name)}
-            )
-            model_columns = {col.name for col in table.c}
+            logger.debug(f"正在检查表 '{table_name}' 的列和索引...")
 
-            missing_columns = model_columns - db_columns
-            if not missing_columns:
-                logger.debug(f"表 '{table_name}' 结构一致，无需修改。")
+            try:
+                # 检查并添加缺失的列
+                db_columns = await connection.run_sync(
+                    lambda conn: {col["name"] for col in inspector.get_columns(table_name, conn)}
+                )
+                model_columns = {col.name for col in table.c}
+                missing_columns = model_columns - db_columns
+
+                if missing_columns:
+                    logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
+                    async with connection.begin() as trans:
+                        for column_name in missing_columns:
+                            try:
+                                column = table.c[column_name]
+                                add_column_ddl = AddColumn(table_name, column)
+                                await connection.execute(add_column_ddl)
+                                logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
+                            except Exception as e:
+                                logger.error(
+                                    f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}",
+                                    exc_info=True,
+                                )
+                                await trans.rollback()
+                                break  # 如果一列失败，则停止处理此表的其他列
+                else:
+                    logger.info(f"表 '{table_name}' 的列结构一致。")
+
+                # 检查并创建缺失的索引
+                db_indexes = await connection.run_sync(
+                    lambda conn: {idx["name"] for idx in inspector.get_indexes(table_name, conn)}
+                )
+                model_indexes = {idx.name for idx in table.indexes}
+                missing_indexes = model_indexes - db_indexes
+
+                if missing_indexes:
+                    logger.info(f"在表 '{table_name}' 中发现缺失的索引: {', '.join(missing_indexes)}")
+                    async with connection.begin() as trans:
+                        for index_name in missing_indexes:
+                            try:
+                                index_obj = next((idx for idx in table.indexes if idx.name == index_name), None)
+                                if index_obj is not None:
+                                    await connection.execute(CreateIndex(index_obj))
+                                    logger.info(f"成功为表 '{table_name}' 创建索引 '{index_name}'。")
+                            except Exception as e:
+                                logger.error(
+                                    f"为表 '{table_name}' 创建索引 '{index_name}' 失败: {e}",
+                                    exc_info=True,
+                                )
+                                await trans.rollback()
+                                break  # 如果一个索引失败，则停止处理此表的其他索引
+                else:
+                    logger.debug(f"表 '{table_name}' 的索引一致。")
+
+            except Exception as e:
+                logger.error(f"在处理表 '{table_name}' 时发生意外错误: {e}", exc_info=True)
                 continue
 
-            logger.info(f"在表 '{table_name}' 中发现缺失的列: {', '.join(missing_columns)}")
-            
-            # 开始事务来添加缺失的列
-            async with connection.begin() as trans:
-                try:
-                    for column_name in missing_columns:
-                        column = table.c[column_name]
-
-                        # 构造并执行 ALTER TABLE 语句
-                        try:
-                            # 在同步上下文中编译列类型
-                            column_type = await connection.run_sync(
-                                lambda sync_conn: column.type.compile(sync_conn.dialect)
-                            )
-                            sql = f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}"
-
-                            # 添加默认值和非空约束的处理
-                            if column.default is not None:
-                                default_value = column.default.arg
-                                if isinstance(default_value, str):
-                                    sql += f" DEFAULT '{default_value}'"
-                                else:
-                                    sql += f" DEFAULT {default_value}"
-
-                            if not column.nullable:
-                                sql += " NOT NULL"
-
-                            await connection.execute(text(sql))
-                            logger.info(f"成功向表 '{table_name}' 添加列 '{column_name}'。")
-                        except Exception as e:
-                            logger.error(f"向表 '{table_name}' 添加列 '{column_name}' 失败: {e}")
-
-                except Exception as e:
-                    logger.error(f"在表 '{table_name}' 添加列时发生错误，事务已回滚: {e}")
-                    await trans.rollback()
-                    raise
-
     logger.info("数据库结构检查与自动迁移完成。")