Py Alembic Patterns

ALWAYS review autogenerated migrations. They often need fixes.

What Autogenerate Catches

• Table creation/deletion
• Column addition/removal
• Column type changes
• Foreign key changes
• Index changes (sometimes)

What Autogenerate Misses

• Column renames (sees as drop + add = DATA LOSS)
• Table renames (same problem)
• Data migrations
• Constraint names
• Partial indexes
• Complex index changes
• Check constraints
• Triggers and functions

# ❌ DANGEROUS: Autogenerated for column rename
def upgrade():
    op.drop_column("users", "name")      # DATA LOSS!
    op.add_column("users", sa.Column("full_name", sa.String()))

# ✅ CORRECT: Manual rename
def upgrade():
    op.alter_column("users", "name", new_column_name="full_name")

def downgrade():
    op.alter_column("users", "full_name", new_column_name="name")

Pattern: Safe Migration Structure

"""Add user preferences table.

Revision ID: abc123
Revises: def456
Create Date: 2024-01-15 10:30:00.000000
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers
revision = "abc123"
down_revision = "def456"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Always explicit, never rely on defaults
    op.create_table(
        "user_preferences",
        sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
        sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column("theme", sa.String(50), nullable=False, server_default="light"),
        sa.Column("notifications_enabled", sa.Boolean(), nullable=False, server_default="true"),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )
    
    # Explicit index names
    op.create_index(
        "ix_user_preferences_user_id",
        "user_preferences",
        ["user_id"],
    )
    
    # Foreign key with explicit name
    op.create_foreign_key(
        "fk_user_preferences_user_id",
        "user_preferences",
        "users",
        ["user_id"],
        ["id"],
        ondelete="CASCADE",
    )


def downgrade() -> None:
    # Always implement downgrade!
    op.drop_constraint("fk_user_preferences_user_id", "user_preferences", type_="foreignkey")
    op.drop_index("ix_user_preferences_user_id", "user_preferences")
    op.drop_table("user_preferences")

Pattern: Adding Non-Nullable Columns

Problem: Adding NOT NULL column to existing table fails if table has rows.

# ❌ WRONG: Fails if table has data
def upgrade():
    op.add_column("users", sa.Column("role", sa.String(50), nullable=False))

# ✅ CORRECT: Three-step process
def upgrade():
    # Step 1: Add as nullable
    op.add_column("users", sa.Column("role", sa.String(50), nullable=True))
    
    # Step 2: Backfill existing rows
    op.execute("UPDATE users SET role = 'member' WHERE role IS NULL")
    
    # Step 3: Add NOT NULL constraint
    op.alter_column("users", "role", nullable=False)

def downgrade():
    op.drop_column("users", "role")

Pattern: Data Migrations

Problem: Need to transform existing data during schema change.

from sqlalchemy import text

def upgrade():
    # Get connection for data operations
    connection = op.get_bind()
    
    # Add new column
    op.add_column("assessments", sa.Column("status", sa.String(20)))
    
    # Migrate data
    connection.execute(
        text("""
            UPDATE assessments 
            SET status = CASE 
                WHEN completed_at IS NOT NULL THEN 'completed'
                WHEN started_at IS NOT NULL THEN 'in_progress'
                ELSE 'pending'
            END
        """)
    )
    
    # Now safe to add NOT NULL
    op.alter_column("assessments", "status", nullable=False)


def downgrade():
    op.drop_column("assessments", "status")

Pattern: Large Table Migrations

Problem: Migrations on large tables can lock the table for too long.

def upgrade():
    # ✅ CORRECT: Add index concurrently (no lock)
    op.execute(
        "CREATE INDEX CONCURRENTLY ix_events_user_id ON events (user_id)"
    )
    
    # Note: CONCURRENTLY requires autocommit mode
    # Add to migration file:
    # from alembic import context
    # context.configure(transaction_per_migration=False)

def downgrade():
    op.execute("DROP INDEX CONCURRENTLY IF EXISTS ix_events_user_id")


# For column changes on large tables, consider:
# 1. Add new column (nullable)
# 2. Backfill in batches via separate script
# 3. Add constraint in separate migration

Pattern: Enum Changes

Problem: PostgreSQL enums are tricky to modify.

# Adding a value to existing enum
def upgrade():
    # PostgreSQL-specific: Add value to enum
    op.execute("ALTER TYPE assessment_status ADD VALUE 'archived'")

def downgrade():
    # Can't remove enum values in PostgreSQL!
    # Options:
    # 1. Leave it (usually fine)
    # 2. Recreate enum (complex, requires data migration)
    pass


# Creating new enum
def upgrade():
    # Create enum type first
    assessment_status = postgresql.ENUM(
        "draft", "active", "completed", "archived",
        name="assessment_status",
        create_type=True,
    )
    assessment_status.create(op.get_bind())
    
    # Then use it
    op.add_column(
        "assessments",
        sa.Column("status", assessment_status, nullable=False, server_default="draft"),
    )

def downgrade():
    op.drop_column("assessments", "status")
    op.execute("DROP TYPE assessment_status")

Pattern: Multiple Heads (Branching)

Problem: Multiple developers creating migrations simultaneously.

# Check for multiple heads
uv run alembic heads

# If multiple heads, create merge migration
uv run alembic merge -m "Merge heads" abc123 def456

# Or specify down_revision as tuple
down_revision = ("abc123", "def456")

Pattern: Testing Migrations

# test_migrations.py
import pytest
from alembic import command
from alembic.config import Config

@pytest.fixture
def alembic_config():
    config = Config("alembic.ini")
    return config

def test_upgrade_downgrade(alembic_config, test_db):
    """Test migrations can upgrade and downgrade."""
    # Upgrade to head
    command.upgrade(alembic_config, "head")
    
    # Downgrade to base
    command.downgrade(alembic_config, "base")
    
    # Upgrade again
    command.upgrade(alembic_config, "head")

def test_migration_has_downgrade():
    """Ensure all migrations have downgrade."""
    # Parse migration files and check downgrade isn't just 'pass'
    ...

Migration Review Checklist

Before applying any migration:

• Downgrade function implemented (not just pass)
• Column renames use alter_column, not drop+add
• Non-nullable columns added with default or backfill
• Large table operations consider locking
• Indexes have explicit names
• Foreign keys have explicit names and ON DELETE behavior
• Enums created before use
• Data migrations tested with real data volumes
• Migration tested: upgrade, downgrade, upgrade

Production Safety

# Set statement timeout to prevent long locks
def upgrade():
    op.execute("SET statement_timeout = '5s'")
    
    # Your migration here
    
    op.execute("SET statement_timeout = '0'")  # Reset

# Always backup before production migrations
pg_dump -h host -U user -d dbname > backup_before_migration.sql

# Apply with --sql to preview
uv run alembic upgrade head --sql

# Apply for real
uv run alembic upgrade head

Common Issues