使用SQLAlchemy与Alembic实现PostgreSQL动态分区的方案问询

完全可以通过Alembic编写迁移脚本，实现从单表到逐月新增时间范围分区的PostgreSQL分区表架构，以下是具体实现步骤和代码示例：

一、初始迁移：将现有单表转换为分区父表 #

如果已有一张普通单表（如events），需要先将其改造为分区父表，并创建第一个包含现有数据的分区：

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
from datetime import datetime, timedelta

def upgrade():
    # 1. 将原表重命名为临时表
    op.rename_table('events', 'events_old')
    
    # 2. 创建分区父表，结构与原表完全一致
    op.create_table(
        'events',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('event_time', postgresql.TIMESTAMP(timezone=True), nullable=False),
        sa.Column('data', sa.JSON(), nullable=True),
        sa.PrimaryKeyConstraint('id')
    )
    
    # 3. 为父表添加分区约束并指定分区方式（按event_time范围分区）
    op.execute("""
        ALTER TABLE events 
        ADD CONSTRAINT events_event_time_not_null CHECK (event_time IS NOT NULL),
        PARTITION BY RANGE (event_time)
    """)
    
    # 4. 创建初始分区，覆盖原表数据的时间范围（这里取当前月份作为初始范围）
    current_month_start = datetime.now().replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    next_month_start = (current_month_start + timedelta(days=32)).replace(day=1)
    
    initial_partition = f"events_{current_month_start.year}_{current_month_start.month:02d}"
    op.execute(f"""
        CREATE TABLE {initial_partition}
        PARTITION OF events 
        FOR VALUES FROM ('{current_month_start.isoformat()}') TO ('{next_month_start.isoformat()}')
    """)
    
    # 5. 将临时表的数据迁移到新分区
    op.execute("INSERT INTO events SELECT * FROM events_old")
    
    # 6. 删除临时表
    op.drop_table('events_old')

def downgrade():
    # 回滚：删除所有分区和父表
    op.execute("DROP TABLE IF EXISTS events_* CASCADE")
    op.drop_table('events')
    # 若需恢复原表，可重新创建表结构并导入备份数据

二、月度分区自动创建脚本 #

编写可重复执行的迁移脚本，每次运行时自动检查并创建下一个月的分区：

from alembic import op
from datetime import datetime, timedelta

def get_next_month_boundaries():
    # 准确计算下一个月的起始和结束时间（避免月末日期溢出问题）
    today = datetime.now()
    current_month_end = (today.replace(day=1) + timedelta(days=32)).replace(day=1)
    next_month_end = (current_month_end + timedelta(days=32)).replace(day=1)
    return current_month_end, next_month_end

def upgrade():
    next_month_start, next_month_end = get_next_month_boundaries()
    partition_name = f"events_{next_month_start.year}_{next_month_start.month:02d}"
    
    # 检查分区是否已存在，避免重复创建
    partition_exists = op.execute(f"""
        SELECT EXISTS (
            SELECT 1 FROM pg_tables 
            WHERE schemaname = 'public' AND tablename = '{partition_name}'
        )
    """).scalar()
    
    if not partition_exists:
        # 创建新分区
        op.execute(f"""
            CREATE TABLE {partition_name}
            PARTITION OF events 
            FOR VALUES FROM ('{next_month_start.isoformat()}') TO ('{next_month_end.isoformat()}')
        """)
        
        # 可选：为分区创建自定义索引（父表的主键/索引会自动继承，无需重复创建）
        # op.execute(f"CREATE INDEX idx_{partition_name}_event_time ON {partition_name} (event_time)")

def downgrade():
    # 回滚：删除已创建的下一个月分区
    next_month_start, _ = get_next_month_boundaries()
    partition_name = f"events_{next_month_start.year}_{next_month_start.month:02d}"
    op.execute(f"DROP TABLE IF EXISTS {partition_name}")

关键注意事项 #

• 数据迁移锁问题：如果原表数据量极大，直接INSERT INTO会导致长时间锁表，建议分批迁移或使用pg_dump导出后导入分区。
• 时区一致性：务必使用带时区的时间类型（TIMESTAMP WITH TIME ZONE），避免时区偏移导致数据落入错误分区。
• 自动化执行：结合crontab或调度工具，每月自动运行alembic upgrade head，即可实现分区的自动创建。
• 旧分区清理：可在迁移脚本中扩展逻辑，自动删除指定时间之前的旧分区（如删除6个月前的分区）。

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