greater_wms/migrate_sqlite_to_postgres.py

import sqlite3
import psycopg2
import os
import sys
from dotenv import load_dotenv
from psycopg2 import sql
from collections import defaultdict

# 加载环境变量
load_dotenv()

class SQLiteToPostgresMigrator:
    def __init__(self):
        # SQLite 配置
        self.sqlite_db_path = os.getenv('SQLITE_DB_PATH', 'db.sqlite3')
        
        # PostgreSQL 配置
        self.pg_config = {
            'dbname': os.getenv('PG_DB_NAME', 'wmsdb'),
            'user': os.getenv('PG_DB_USER', 'wmsuser'),
            'password': os.getenv('PG_DB_PASSWORD', 'abc@1234'),
            'host': os.getenv('PG_DB_HOST', 'localhost'),
            'port': os.getenv('PG_DB_PORT', '5432')
        }
        
        # 连接
        self.sqlite_conn = None
        self.sqlite_cursor = None
        self.pg_conn = None
        self.pg_cursor = None
        
        # 表依赖关系
        self.table_dependencies = defaultdict(list)
        self.reverse_dependencies = defaultdict(list)
        self.migration_order = []
        
        # 需要跳过的系统表
        self.skip_tables = [
            'auth_permission',
            'django_content_type',
            'django_migrations',
            'auth_group_permissions',
            'auth_user_groups',
            'auth_user_user_permissions'
        ]
        
        # 迁移状态
        self.migrated_tables = set()
        self.failed_tables = set()
        
        # 存储表的主键信息
        self.primary_keys = {}
        
        # 存储字段长度限制
        self.column_lengths = {}

    def connect(self):
        """连接到数据库"""
        try:
            # SQLite 连接
            self.sqlite_conn = sqlite3.connect(self.sqlite_db_path)
            self.sqlite_cursor = self.sqlite_conn.cursor()
            
            # PostgreSQL 连接
            self.pg_conn = psycopg2.connect(**self.pg_config)
            self.pg_cursor = self.pg_conn.cursor()
            print(f"连接成功: SQLite({self.sqlite_db_path}) → PostgreSQL({self.pg_config['dbname']})")
            return True
        except Exception as e:
            print(f"连接失败: {str(e)}")
            return False

    def close(self):
        """关闭数据库连接"""
        for cursor in [self.sqlite_cursor, self.pg_cursor]:
            if cursor: cursor.close()
        for conn in [self.sqlite_conn, self.pg_conn]:
            if conn: conn.close()

    def clear_postgres_data(self):
        """清空 PostgreSQL 数据库中的所有数据"""
        print("\n警告: 此操作将清空 PostgreSQL 数据库中的所有数据!")
        confirm = input("确定要清空 PostgreSQL 数据库吗? (y/n): ")
        
        if confirm.lower() != 'y':
            print("取消清空操作")
            return False
        
        try:
            # 获取所有表名
            self.pg_cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public'
            """)
            tables = [row[0] for row in self.pg_cursor.fetchall()]
            
            # 使用 TRUNCATE TABLE ... CASCADE 替代 session_replication_role
            for table in tables:
                try:
                    self.pg_cursor.execute(f"TRUNCATE TABLE {table} CASCADE;")
                    print(f"已清空表: {table}")
                except Exception as e:
                    print(f"清空表 {table} 时出错: {str(e)}")
                    continue
            
            self.pg_conn.commit()
            print("成功清空 PostgreSQL 数据库中的所有数据")
            return True
        except Exception as e:
            self.pg_conn.rollback()
            print(f"清空数据库时出错: {str(e)}")
            return False

    def analyze_dependencies(self):
        """分析表依赖关系"""
        # 获取 SQLite 中的所有表
        self.sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = [row[0] for row in self.sqlite_cursor.fetchall()]
        
        # 排除系统表
        excluded_tables = ['sqlite_sequence'] + self.skip_tables
        tables = [t for t in tables if t not in excluded_tables]
        
        # 分析外键依赖
        for table in tables:
            self.sqlite_cursor.execute(f"PRAGMA foreign_key_list({table})")
            for fk in self.sqlite_cursor.fetchall():
                ref_table = fk[2]
                if ref_table in tables:
                    self.table_dependencies[ref_table].append(table)
                    self.reverse_dependencies[table].append(ref_table)
        
        # 创建迁移顺序 (拓扑排序)
        visited = set()
        order = []
        
        def visit(table):
            if table in visited:
                return
            visited.add(table)
            for dependent in self.table_dependencies.get(table, []):
                visit(dependent)
            order.append(table)
        
        for table in tables:
            if table not in visited:
                visit(table)
        
        self.migration_order = order
        print(f"分析完成: 发现 {len(tables)} 个表, 依赖顺序已确定")

    def get_pg_table_structure(self, table_name):
        """获取 PostgreSQL 表结构"""
        try:
            self.pg_cursor.execute("""
                SELECT column_name, data_type, character_maximum_length 
                FROM information_schema.columns 
                WHERE table_name = %s
            """, (table_name.lower(),))
            
            structure = {}
            for row in self.pg_cursor.fetchall():
                column_name, data_type, max_length = row
                structure[column_name] = {
                    'data_type': data_type,
                    'max_length': max_length
                }
            return structure
        except Exception:
            return {}

    def get_column_lengths(self, table_name):
        """获取 PostgreSQL 表的字段长度限制"""
        if table_name in self.column_lengths:
            return self.column_lengths[table_name]
        
        try:
            self.pg_cursor.execute("""
                SELECT column_name, character_maximum_length 
                FROM information_schema.columns 
                WHERE table_name = %s
            """, (table_name.lower(),))
            
            lengths = {}
            for row in self.pg_cursor.fetchall():
                column_name, max_length = row
                if max_length:
                    lengths[column_name] = max_length
            self.column_lengths[table_name] = lengths
            return lengths
        except Exception:
            return {}

    def transform_data(self, row, columns, pg_structure):
        """根据目标类型转换数据"""
        transformed = []
        for i, value in enumerate(row):
            col_name = columns[i]
            col_info = pg_structure.get(col_name.lower(), {})
            pg_type = col_info.get('data_type')
            max_length = col_info.get('max_length')
            
            # 处理布尔字段转换
            if pg_type in ('boolean', 'bool'):
                # 将整数0/1转换为布尔值
                if value in (0, '0', False):
                    transformed.append(False)
                elif value in (1, '1', True):
                    transformed.append(True)
                else:
                    # 无法转换的值保持原样（可能会出错）
                    transformed.append(value)
            # 处理字符串长度限制
            elif pg_type in ('character varying', 'varchar', 'char', 'text') and max_length and isinstance(value, str):
                if len(value) > max_length:
                    # 截断超过长度的字符串
                    truncated = value[:max_length]
                    print(f"  警告: 字段 '{col_name}' 值被截断 ({len(value)} -> {max_length})")
                    transformed.append(truncated)
                else:
                    transformed.append(value)
            else:
                transformed.append(value)
        return transformed

    def get_primary_keys(self, table):
        """获取表的主键列"""
        if table in self.primary_keys:
            return self.primary_keys[table]
        
        # 查询SQLite表的主键
        self.sqlite_cursor.execute(f"PRAGMA table_info({table})")
        columns_info = self.sqlite_cursor.fetchall()
        
        # 主键列是那些pk值不为0的列
        primary_keys = []
        for col_info in columns_info:
            # col_info格式: (cid, name, type, notnull, dflt_value, pk)
            if col_info[5] != 0:  # pk列不为0表示是主键
                primary_keys.append(col_info[1])
        
        self.primary_keys[table] = primary_keys
        return primary_keys

    def check_dependencies_migrated(self, table):
        """检查依赖的表是否已经成功迁移"""
        dependencies = self.reverse_dependencies.get(table, [])
        
        if not dependencies:
            return True
        
        print(f"  检查 {table} 的依赖表迁移状态...")
        all_migrated = True
        
        for dep_table in dependencies:
            if dep_table in self.migrated_tables:
                print(f"    ✓ {dep_table} 已迁移")
            elif dep_table in self.failed_tables:
                print(f"    ✗ {dep_table} 迁移失败")
                all_migrated = False
            else:
                print(f"    ? {dep_table} 尚未迁移")
                all_migrated = False
        
        return all_migrated

    def migrate_table(self, table):
        """迁移单个表"""
        print(f"\n准备迁移表: {table}")
        
        # 检查是否需要跳过
        if table in self.skip_tables:
            print(f"  跳过系统表: {table}")
            return True
        
        # 检查依赖表是否已迁移
        if not self.check_dependencies_migrated(table):
            print(f"  依赖表未完全迁移，暂时跳过 {table}")
            return False
        
        # 获取 SQLite 表结构
        self.sqlite_cursor.execute(f"PRAGMA table_info({table})")
        columns_info = self.sqlite_cursor.fetchall()
        columns = [col[1] for col in columns_info]
        
        # 获取主键
        primary_keys = self.get_primary_keys(table)
        
        # 获取数据
        self.sqlite_cursor.execute(f"SELECT * FROM {table}")
        rows = self.sqlite_cursor.fetchall()
        
        if not rows:
            print(f"  表 {table} 为空，跳过")
            self.migrated_tables.add(table)
            return True
        
        # 获取 PostgreSQL 表结构
        pg_structure = self.get_pg_table_structure(table)
        
        # 获取字段长度限制
        column_lengths = self.get_column_lengths(table)
        
        # 创建 PostgreSQL 插入语句
        placeholders = ', '.join(['%s'] * len(columns))
        column_names = ', '.join([f'"{col}"' for col in columns])
        
        # 构建插入语句，根据主键是否存在添加ON CONFLICT子句
        if primary_keys:
            # 使用主键处理冲突
            conflict_columns = ', '.join([f'"{pk}"' for pk in primary_keys])
            insert_query = sql.SQL("""
                INSERT INTO {} ({}) 
                VALUES ({})
                ON CONFLICT ({}) DO NOTHING
            """).format(
                sql.Identifier(table.lower()),
                sql.SQL(column_names),
                sql.SQL(placeholders),
                sql.SQL(conflict_columns)
            )
        else:
            # 没有主键，使用普通插入语句
            insert_query = sql.SQL("""
                INSERT INTO {} ({}) 
                VALUES ({})
            """).format(
                sql.Identifier(table.lower()),
                sql.SQL(column_names),
                sql.SQL(placeholders)
            )
        
        # 用户确认
        print(f"  表 {table} 有 {len(rows)} 行数据需要迁移")
        if primary_keys:
            print(f"  使用主键 ({', '.join(primary_keys)}) 处理冲突")
        else:
            print("  警告: 表没有主键，可能产生重复数据")
        
        # 显示字段长度限制
        if column_lengths:
            print("  字段长度限制:")
            for col, max_len in column_lengths.items():
                print(f"    {col}: {max_len} 字符")
        
        confirm = input("  按 Enter 键开始迁移，或输入 's' 跳过此表: ")
        
        if confirm.lower() == 's':
            print(f"  已跳过表 {table}")
            return True
        
        # 迁移数据
        try:
            success_count = 0
            error_count = 0
            for row in rows:
                try:
                    # 转换数据类型
                    transformed_row = self.transform_data(row, columns, pg_structure)
                    self.pg_cursor.execute(insert_query, transformed_row)
                    success_count += 1
                except psycopg2.IntegrityError as e:
                    error_count += 1
                    if 'foreign key constraint' in str(e):
                        print(f"  外键约束错误: {str(e)}")
                    else:
                        print(f"  完整性错误: {str(e)}")
                    self.pg_conn.rollback()
                except Exception as e:
                    error_count += 1
                    print(f"  行插入失败: {str(e)}")
                    self.pg_conn.rollback()
            
            self.pg_conn.commit()
            
            if error_count == 0:
                print(f"  成功迁移 {success_count} 行数据")
                self.migrated_tables.add(table)
                return True
            else:
                print(f"  部分迁移: {success_count} 行成功, {error_count} 行失败")
                self.failed_tables.add(table)
                return False
        except Exception as e:
            self.pg_conn.rollback()
            print(f"  迁移表 {table} 时出错: {str(e)}")
            self.failed_tables.add(table)
            return False

    def migrate(self):
        """执行迁移"""
        if not self.connect():
            return False
        
        try:
            # 询问是否清空 PostgreSQL 数据
            print("\n是否在迁移前清空 PostgreSQL 数据库?")
            print("1. 清空所有数据 (确保无重复)")
            print("2. 不清空 (可能产生重复数据)")
            print("3. 退出")
            
            choice = input("请选择 (1/2/3): ")
            
            if choice == '3':
                print("退出迁移")
                return False
                
            if choice == '1':
                if not self.clear_postgres_data():
                    print("清空操作失败，退出迁移")
                    return False
            
            self.analyze_dependencies()
            print(f"\n开始迁移 {len(self.migration_order)} 个表")
            
            # 按依赖顺序迁移表
            remaining_tables = self.migration_order.copy()
            max_attempts = 10  # 最大尝试次数
            attempt = 0
            
            while remaining_tables and attempt < max_attempts:
                attempt += 1
                print(f"\n第 {attempt} 轮迁移尝试")
                
                # 复制列表以便在迭代时修改
                tables_to_migrate = remaining_tables.copy()
                remaining_tables = []
                
                for table in tables_to_migrate:
                    success = self.migrate_table(table)
                    if not success:
                        remaining_tables.append(table)
                
                if remaining_tables:
                    print(f"\n本轮迁移后仍有 {len(remaining_tables)} 个表未成功迁移")
            
            # 检查迁移结果
            if remaining_tables:
                print(f"\n迁移完成，但以下表迁移失败:")
                for table in remaining_tables:
                    print(f"  - {table}")
            else:
                print("\n所有表迁移成功!")
            
            return True
        except Exception as e:
            print(f"迁移过程中发生错误: {str(e)}")
            return False
        finally:
            self.close()

if __name__ == "__main__":
    migrator = SQLiteToPostgresMigrator()
    migrator.migrate()