Database migrations with Alembic

Motivation

Your AI backend ships, users store conversations, and a week later you need to add a cost_usd column to the messages table. How do you add that column to the live production database without deleting the existing data?

Base.metadata.create_all() from Lesson 22 cannot help — it only creates tables that don’t exist yet. Manually running ALTER TABLE SQL against production is dangerous and untracked. Database migrations are the answer: version-controlled, repeatable, reversible SQL changes that you can apply in order on any environment. Alembic is the migration tool for SQLAlchemy — it is the standard for Python AI backends, and it is expected knowledge in AI engineering roles that touch production databases.

What migrations are and why they exist

A migration is a Python script that describes one specific change to a database schema — adding a column, dropping a table, changing a column type. Each migration has two functions:

• upgrade() — applies the change (e.g. adds the column)
• downgrade() — reverses it (e.g. drops the column)

Alembic keeps track of which migrations have been applied in a special table in your database (alembic_version). When you run alembic upgrade head, it applies only the migrations that haven’t been applied yet — in the correct order.

This gives you three things development-only tools like create_all cannot provide:

1. Safety — you can test on staging before touching production
2. Reversibility — you can roll back to the previous state if something breaks
3. History — every schema change is code-reviewed alongside the application code that uses it

Setting up Alembic

uv add alembic

Initialise Alembic in your project:

alembic init alembic

This creates an alembic/ directory and an alembic.ini file:

your-project/
├── alembic/
│   ├── env.py              # migration environment configuration
│   ├── versions/           # individual migration files go here
│   └── script.py.mako      # template for new migration files
├── alembic.ini             # Alembic configuration
├── main.py
└── database.py

Configure the database URL

Open alembic.ini and find the sqlalchemy.url line. Replace it with a reference to an environment variable (never hardcode credentials in a committed file):

# alembic.ini
sqlalchemy.url = %(DATABASE_URL)s

Then in alembic/env.py, read the variable from the environment:

import os
from dotenv import load_dotenv

load_dotenv()

# Override the sqlalchemy.url with the env var
config.set_main_option("sqlalchemy.url",
    os.getenv("DATABASE_URL", "postgresql+psycopg2://postgres:password@localhost:5432/aidb")
)

Point Alembic at your models

Alembic’s autogenerate feature compares your SQLAlchemy models against the live database and generates the SQL to make them match. To enable this, open alembic/env.py and add your models’ metadata:

from database import Base   # your SQLAlchemy Base

# ... existing env.py code ...
target_metadata = Base.metadata  # ← replace "target_metadata = None" with this

Creating and running migrations

Initial migration — create all tables

The first migration creates all your tables. Run autogenerate, which compares your models to the empty database and generates the SQL:

# Generate a migration script called "create_initial_tables"
alembic revision --autogenerate -m "create_initial_tables"

This creates a file in alembic/versions/ like abc123def456_create_initial_tables.py. Open it and read what it generated — always review autogenerated migrations before running them:

"""create initial tables

Revision ID: abc123def456
Revises:
Create Date: 2026-01-15 14:30:00.000000
"""
from alembic import op
import sqlalchemy as sa

def upgrade() -> None:
    op.create_table(
        'conversations',
        sa.Column('id',         sa.Integer(), autoincrement=True, nullable=False),
        sa.Column('title',      sa.String(255), nullable=False),
        sa.Column('model',      sa.String(100), nullable=False),
        sa.Column('created_at', sa.DateTime(), server_default=sa.func.now(), nullable=True),
        sa.PrimaryKeyConstraint('id'),
    )
    op.create_table(
        'messages',
        sa.Column('id',              sa.Integer(), autoincrement=True, nullable=False),
        sa.Column('conversation_id', sa.Integer(), nullable=False),
        sa.Column('role',            sa.String(20), nullable=False),
        sa.Column('content',         sa.Text(), nullable=False),
        sa.Column('tokens_used',     sa.Integer(), nullable=True),
        sa.Column('created_at',      sa.DateTime(), server_default=sa.func.now(), nullable=True),
        sa.ForeignKeyConstraint(['conversation_id'], ['conversations.id'], ondelete='CASCADE'),
        sa.PrimaryKeyConstraint('id'),
    )

def downgrade() -> None:
    op.drop_table('messages')
    op.drop_table('conversations')

Apply the migration:

alembic upgrade head

head means “apply all unapplied migrations up to the latest one”. Check the database — the tables now exist.

Adding a new column

Two weeks later, you want to track the cost in USD of each message. Add the column to your SQLAlchemy model first:

class Message(Base):
    # ... existing columns ...
    cost_usd: Mapped[float | None] = mapped_column(Float, nullable=True)

Then generate a new migration:

alembic revision --autogenerate -m "add_cost_usd_to_messages"

Alembic compares the model definition against the current database schema and generates only the diff — in this case, an ALTER TABLE:

def upgrade() -> None:
    op.add_column('messages', sa.Column('cost_usd', sa.Float(), nullable=True))

def downgrade() -> None:
    op.drop_column('messages', 'cost_usd')

Apply it:

alembic upgrade head

Essential Alembic commands

The standard migration workflow

Once Alembic is set up, the workflow for every schema change is the same:

1. Change the SQLAlchemy model — add the column, table, or index
2. Generate the migration — alembic revision --autogenerate -m "description"
3. Review the generated file — autogenerate is good but not perfect; always read it
4. Test locally — alembic upgrade head against your dev database
5. Commit both files — the model change and the migration file in the same commit
6. Apply in production — alembic upgrade head as part of your deploy process

Running migrations in Docker

In a Dockerised setup, you run migrations before starting the application. The simplest approach is an entrypoint script:

#!/bin/bash
# entrypoint.sh
set -e

echo "Running database migrations..."
alembic upgrade head

echo "Starting application..."
exec uvicorn main:app --host 0.0.0.0 --port 8000

Update the Dockerfile to use it:

COPY entrypoint.sh .
RUN chmod +x entrypoint.sh
CMD ["./entrypoint.sh"]

Now every time the container starts, it first applies any pending migrations, then starts the server. This is safe because Alembic’s migration tracking is idempotent — running upgrade head when already at head does nothing.