From 4d8fe412d2fbcc4f67dedf1e9285b9d0164667f3 Mon Sep 17 00:00:00 2001
From: hanxuanyu <2252193204@qq.com>
Date: Thu, 28 Mar 2024 12:44:51 +0800
Subject: [PATCH] =?UTF-8?q?=E6=AF=8F=E6=AC=A1=E5=A4=8D=E5=88=B6=E5=89=8D?=
 =?UTF-8?q?=E5=88=A0=E9=99=A4=E6=89=80=E6=9C=89=E8=A1=A8=E5=B9=B6=E9=87=8D?=
 =?UTF-8?q?=E6=96=B0=E5=88=9B=E5=BB=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 database_migrate/db_migrate.py | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

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):