Writing data migrations safely in Alembic

Write the backfill against a lightweight table definition declared inside the revision, run it in bounded batches that each commit, and keep it in a separate revision from the DDL — importing your application models into a migration ties a historical revision to code that will change underneath it. This is the data half of managing multiple heads and database branching.

Quick Answer

Before — a data migration that imports models and does it all at once:

from myapp.models import Order          # couples this revision to current code


def upgrade() -> None:
    session = Session(bind=op.get_bind())
    for order in session.query(Order).all():        # loads the whole table
        order.status_code = STATUS_MAP[order.status]
    session.commit()                                 # one enormous transaction

After — a self-contained, batched backfill:

from __future__ import annotations

import sqlalchemy as sa
from alembic import op

revision = "9c14ab02f7de"
down_revision = "8f2c14a9b7d3"

# Declared here, not imported: this describes the table as it is AT THIS REVISION.
orders = sa.table(
    "orders",
    sa.column("id", sa.Integer),
    sa.column("status", sa.String),
    sa.column("status_code", sa.SmallInteger),
)

STATUS_CODES = {"new": 1, "paid": 2, "shipped": 3, "cancelled": 9}
BATCH = 5_000


def upgrade() -> None:
    connection = op.get_bind()
    last_id = 0
    while True:
        ids = connection.execute(
            sa.select(orders.c.id)
            .where(orders.c.id > last_id, orders.c.status_code.is_(None))
            .order_by(orders.c.id)
            .limit(BATCH)
        ).scalars().all()
        if not ids:
            break
        connection.execute(
            sa.update(orders)
            .where(orders.c.id.in_(ids))
            .values(status_code=sa.case(STATUS_CODES, value=orders.c.status, else_=0))
        )
        last_id = ids[-1]


def downgrade() -> None:
    # Data migrations are frequently irreversible. Say so rather than pretending.
    pass

The sa.table() declaration is the important part: this revision keeps working after the Order model grows three more columns, because it never mentions the model at all.

Execution Context & Async Workflow Integration

A migration runs through op.get_bind(), which is a synchronous connection even when env.py drives an async engine — the greenlet bridge established by run_sync makes it look ordinary. That means a data migration is written with plain synchronous SQLAlchemy, and attempting to await anything inside it fails.

Whether each batch commits depends on transaction_per_migration. With the default single transaction for the whole upgrade, the batched loop above still holds every row lock it takes until the end — the batching bounds memory but not locking. Making it genuinely incremental requires committing between batches, which means the backfill belongs outside the migration:

"""backfill_status_codes.py — a resumable job, run after the migration."""
from __future__ import annotations

import asyncio

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from myapp.models import Order          # a JOB may import models; a migration may not

BATCH = 5_000


async def backfill(url: str) -> int:
    engine = create_async_engine(url)
    Session = async_sessionmaker(engine, expire_on_commit=False)
    processed, last_id = 0, 0
    try:
        while True:
            async with Session() as session, session.begin():   # commits per batch
                ids = (await session.execute(
                    select(Order.id)
                    .where(Order.id > last_id, Order.status_code.is_(None))
                    .order_by(Order.id)
                    .limit(BATCH)
                )).scalars().all()
                if not ids:
                    return processed
                await session.execute(
                    update(Order).where(Order.id.in_(ids)).values(
                        status_code=Order.status_code
                    )
                )
                last_id, processed = ids[-1], processed + len(ids)
            await asyncio.sleep(0.05)      # let replication and vacuum keep up
    finally:
        await engine.dispose()

The asyncio.sleep between batches is not padding. A tight backfill loop generates write-ahead log faster than a replica can apply it, and replication lag is what turns a routine backfill into a user-visible incident on read replicas — the read-after-write window described in routing reads to replicas with async engines widens with exactly this kind of load.

A job may import models, because it is deployed alongside them and is deleted once it has run. A migration may not, because it is kept forever.

Resolving Warnings, Errors & Common Mistakes

Advanced: Making a Backfill Observable and Interruptible

A backfill that runs for an hour needs to be watchable, and the two things worth reporting are progress and remaining work.

from __future__ import annotations

import logging
import signal
import time

logger = logging.getLogger("backfill")
_stop = False


def _request_stop(signum, frame) -> None:
    """SIGTERM asks the loop to finish its current batch and exit cleanly."""
    global _stop
    _stop = True


signal.signal(signal.SIGTERM, _request_stop)
signal.signal(signal.SIGINT, _request_stop)


async def backfill_observable(url: str, batch: int = 5_000) -> None:
    engine = create_async_engine(url)
    Session = async_sessionmaker(engine, expire_on_commit=False)

    async with Session() as session:
        remaining = await session.scalar(
            select(func.count()).select_from(Order).where(Order.status_code.is_(None))
        )
    logger.info("backfill starting", extra={"remaining": remaining})

    processed, started, last_id = 0, time.monotonic(), 0
    while not _stop:
        async with Session() as session, session.begin():
            ids = (await session.execute(
                select(Order.id)
                .where(Order.id > last_id, Order.status_code.is_(None))
                .order_by(Order.id).limit(batch)
            )).scalars().all()
            if not ids:
                break
            await session.execute(
                update(Order).where(Order.id.in_(ids)).values(status_code=1)
            )
            last_id, processed = ids[-1], processed + len(ids)

        rate = processed / max(time.monotonic() - started, 1e-9)
        logger.info("backfill progress", extra={
            "processed": processed,
            "remaining": max(remaining - processed, 0),
            "rows_per_second": round(rate, 1),
            "eta_seconds": round(max(remaining - processed, 0) / max(rate, 1e-9)),
        })
        await asyncio.sleep(0.05)

    logger.info("backfill stopped", extra={"processed": processed,
                                           "completed": not _stop})

Two properties matter here. The signal handler makes the job interruptible at a batch boundary, so stopping it during an incident leaves the database consistent and the work resumable — an ordinary kill mid-batch would too, thanks to the per-batch transaction, but a clean stop also logs where it got to. And the estimate is reported continuously rather than computed once, so a backfill that slows down as it reaches a colder part of the table says so.

The predicate status_code IS NULL is what makes the job idempotent: re-running it processes only what remains, so a restart is always safe and a partially completed run is never a problem. Where the target column has no natural "not yet done" value, add a nullable marker column for the duration of the backfill and drop it afterwards — that is cheaper than the alternative of tracking progress in a separate table nobody remembers to clean up.

Splitting DDL From Data

The strongest structural rule is that a revision should either change the schema or move data, and not both.

The reasons compound. A DDL revision is fast and predictable; a data revision is slow and depends on how much data exists, which differs by environment. Combining them means the deploy takes the sum and the fast part cannot be applied without the slow part. A failure part-way leaves an ambiguous state: was the column added? was some of the data moved? Splitting them makes each answer independently checkable with alembic current.

The sequence for adding a derived column therefore has three steps rather than one, exactly matching the expand-and-contract ladder in zero-downtime schema migration strategies:

1. A DDL revision adds the column as nullable with a constant default. Fast, and safe on a live table.
2. A separate job backfills it in batches, resumable and observable, while both old and new code are running.
3. A second DDL revision adds the NOT NULL constraint as NOT VALID and validates it — which will fail if the backfill has not completed, and that failure is exactly the check you want.

That third step is worth appreciating: it makes the constraint the verification. There is no need to remember to check whether the backfill finished, because the migration that depends on it refuses to apply otherwise. Encoding a precondition as a constraint rather than as a runbook step is the difference between a process that works when everyone is paying attention and one that works when nobody is.

Frequently Asked Questions

Related

• Managing Multiple Heads and Database Branching — The parent guide: revision structure and why data migrations get their own revision.
• Running Alembic migrations in CI/CD pipelines — Where a backfill job is triggered relative to the migration.
• Zero-downtime schema migration strategies — The expand-and-contract ladder this three-step sequence belongs to.
• Adding a NOT NULL column without locking in PostgreSQL — The DDL either side of the backfill, and the lock each step takes.