Alembic与SQLAlchemy：非分区表数据迁移至分区表的最佳实践

先修正你现有方法的问题 #

1. COPY语法错误：你写的COPY partitioned_table FROM original_table不符合PostgreSQL语法，COPY FROM只能读取文件或标准输入，不能直接从另一个表取数。正确的跨表COPY需要用管道：

psql -c "COPY original_table TO STDOUT" | psql -c "COPY partitioned_table FROM STDIN"

如果要在Alembic脚本里执行，可以通过调用shell命令实现。

2. INSERT INTO ... SELECT失败的常见原因：
   • 原表数据的分区键值不匹配新分区表的分区规则（比如按日期分区的表，原表存在不在任何子分区范围内的数据）
   • 新分区表的约束（非空、唯一键等）比原表严格，原表有不符合约束的数据
   • 数据量过大导致事务超时、锁表，或者触发数据库的资源限制
   • 新旧表字段类型不兼容（比如原表是timestamp，分区表是timestamptz）

结合Alembic/SQLAlchemy的生产级迁移步骤 #

1. 前期准备 #

• 确保新分区表的结构（字段、类型、约束、索引）和原表完全一致，仅添加postgresql_partition_by参数并创建好对应子分区
• 若业务允许，选择低峰期执行迁移，暂时暂停原表的写入操作；如果不能停写，需提前配置逻辑复制同步增量数据
• 备份原表，避免迁移失败导致数据丢失

2. 数据迁移的高效方案 #

方案一：批量INSERT（中小表适用）

数据量不大时，用分批INSERT避免锁表，同时确保字段完全匹配：

-- 单批次示例，可循环执行
INSERT INTO partitioned_table
SELECT * FROM original_table
WHERE id BETWEEN 1 AND 10000;

在Alembic脚本中实现分批逻辑：

from sqlalchemy import text

def upgrade():
    batch_size = 10000
    offset = 0
    while True:
        # 检查当前批次是否有数据
        count = op.execute(
            text("SELECT COUNT(*) FROM original_table WHERE id BETWEEN :start AND :end"),
            {"start": offset, "end": offset + batch_size}
        ).scalar()
        if count == 0:
            break
        # 执行批量插入
        op.execute(
            text("INSERT INTO partitioned_table SELECT * FROM original_table WHERE id BETWEEN :start AND :end"),
            {"start": offset, "end": offset + batch_size}
        )
        offset += batch_size

方案二：COPY命令（大表首选，性能更高）

COPY的吞吐量远高于INSERT，适合百万级以上数据。在Alembic中通过subprocess调用shell命令：

import subprocess

def upgrade():
    # 替换为你的数据库名称
    db_name = "your_database_name"
    # 执行COPY管道
    subprocess.run(
        ["sh", "-c", f'psql -d {db_name} -c "COPY original_table TO STDOUT" | psql -d {db_name} -c "COPY partitioned_table FROM STDIN"'],
        check=True
    )

注意：要确保执行脚本的用户有psql访问权限，可通过环境变量或pgpass文件配置数据库认证信息。

方案三：pg_dump导出导入（超大规模表）

对于超大型表，用pg_dump单独导出原表再导入分区表更稳妥：

# 导出原表
pg_dump -t original_table your_database_name > original_table_dump.sql
# 导入到分区表
psql your_database_name -c "\copy partitioned_table FROM original_table_dump.sql"

同样可以在Alembic中用subprocess调用这些命令。

3. 数据一致性验证 #

迁移完成后必须验证数据完整：

-- 检查行数是否一致
SELECT COUNT(*) FROM original_table;
SELECT COUNT(*) FROM partitioned_table;

-- 随机抽样对比数据
SELECT * FROM original_table ORDER BY random() LIMIT 10;
SELECT * FROM partitioned_table ORDER BY random() LIMIT 10;

-- 数值型字段求和验证
SELECT SUM(numeric_column) FROM original_table;
SELECT SUM(numeric_column) FROM partitioned_table;

4. 原子切换表名 #

用事务包裹删除和重命名操作，避免业务中断：

BEGIN;
DROP TABLE original_table;
ALTER TABLE partitioned_table RENAME TO original_table;
COMMIT;

在Alembic中实现：

def upgrade():
    # 开启事务
    with op.get_bind().begin():
        op.drop_table("original_table")
        op.rename_table("partitioned_table", "original_table")

5. 后续检查 #

• 恢复原表的写入操作（如果之前暂停）
• 验证分区表的查询路由是否正常，确保查询只扫描对应分区
• 监控分区表的存储、索引性能

内容的提问来源于stack exchange，提问作者Mohammad Almokdad