diff --git a/database_migrate/db_migrate.py b/database_migrate/db_migrate.py index 8e3cc27..e87e973 100644 --- a/database_migrate/db_migrate.py +++ b/database_migrate/db_migrate.py @@ -87,20 +87,37 @@ class DBMigrator: self.batch_size = batch_size def clear_target_database(self): - # 清空目标数据库所有表 + # 彻底删除目标数据库所有表 try: with DatabaseConnection(self.target_config) as target_db: with target_db.cursor() as cursor: + cursor.execute("SET FOREIGN_KEY_CHECKS = 0;") cursor.execute("SHOW TABLES;") tables = cursor.fetchall() for table in tables: - cursor.execute(f"TRUNCATE TABLE {table[0]};") + cursor.execute(f"DROP TABLE IF EXISTS {table[0]};") + cursor.execute("SET FOREIGN_KEY_CHECKS = 1;") target_db.commit() return True except Exception as e: logging.error(f"Error clearing target database: {e}") return False + def copy_table_structure(self, table): + # 复制表结构 + try: + with DatabaseConnection(self.source_config) as source_db, DatabaseConnection( + self.target_config) as target_db: + with source_db.cursor() as source_cursor, target_db.cursor() as target_cursor: + source_cursor.execute(f"SHOW CREATE TABLE {table};") + create_table_sql = source_cursor.fetchone()[1] + target_cursor.execute(create_table_sql) + target_db.commit() + return True + except Exception as e: + logging.error(f"Error copying table structure for {table}: {e}") + return False + def migrate_table_data(self, table): # 为每个表迁移数据 try: @@ -153,6 +170,12 @@ class DBMigrator: def migrate(self, concurrency): self.clear_target_database() tables = self.get_tables() + # 复制所有表结构 + for table in tables: + if not self.copy_table_structure(table): + logging.error(f"Failed to copy structure for table {table}. Migration aborted.") + return + # 使用线程池并发迁移数据 with ThreadPoolExecutor(max_workers=concurrency) as executor: futures = {executor.submit(self.migrate_table_data, table): table for table in tables} for future in as_completed(futures):