FastAPI Database Migrations

In this tutorial, we'll explore database migrations in FastAPI applications. Database migrations are a systematic way to evolve your database schema over time while preserving existing data. They are essential for maintaining and updating applications in production environments.

Introduction to Database Migrations

When developing applications that use databases, the database schema inevitably changes as the application evolves. You might need to:

• Add new tables
• Add columns to existing tables
• Remove or rename columns
• Change data types
• Create new indexes

Making these changes manually is error-prone and difficult to track. Database migrations solve this problem by providing a structured approach to evolve your database schema.

Why Use Database Migrations?

• Version control for your database: Track changes to your database schema just like you track changes to your code.
• Reproducibility: Ensure consistent database setup across development, testing, and production environments.
• Collaboration: Allow multiple developers to make and share database changes.
• Safe production deployments: Apply changes to production databases with confidence.
• Rollback capability: Revert database changes if necessary.

Alembic: The Migration Tool for SQLAlchemy

In the FastAPI ecosystem, most applications use SQLAlchemy as the ORM (Object-Relational Mapper). For database migrations with SQLAlchemy, Alembic is the recommended tool. Alembic was created by the same author as SQLAlchemy and integrates perfectly with it.

Let's see how to set up and use Alembic with FastAPI.

Setting Up Alembic in a FastAPI Project

Step 1: Install Required Packages

bash

pip install alembic sqlalchemy fastapi

Step 2: Initialize Alembic

Create an alembic directory in your project and initialize it:

bash

mkdir alembic
cd alembic
alembic init migrations

This creates an alembic.ini file and a migrations directory.

Step 3: Configure Alembic

Open the alembic.ini file and set the database URL:

ini

# alembic.ini
sqlalchemy.url = postgresql://username:password@localhost/dbname

Next, modify the env.py file in the migrations directory to use your SQLAlchemy models:

python

# migrations/env.py
from alembic import context
from sqlalchemy import engine_from_config, pool

from app.database import Base  # Import your SQLAlchemy Base
from app.config import settings  # Import your app settings

# this is the Alembic Config object
config = context.config

# Set the database URL from your application settings
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)

# Add your models here
# target_metadata = None
target_metadata = Base.metadata

# other code...

Creating and Running Migrations

Step 1: Create a Migration

After making changes to your SQLAlchemy models, create a new migration:

bash

alembic revision --autogenerate -m "Add user table"

This command will:

1. Compare your SQLAlchemy models with the current database state
2. Generate a migration script in the migrations/versions/ directory

The generated script will look something like:

python

# migrations/versions/a1b2c3d4e5f6_add_user_table.py
"""Add user table

Revision ID: a1b2c3d4e5f6
Revises: 
Create Date: 2023-06-15 10:30:45.123456

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers
revision = 'a1b2c3d4e5f6'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    # Create the users table
    op.create_table('users',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('email', sa.String(), nullable=False),
        sa.Column('hashed_password', sa.String(), nullable=False),
        sa.Column('is_active', sa.Boolean(), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_users_email'), 'users', ['email'], unique=True)


def downgrade():
    # Drop the users table
    op.drop_index(op.f('ix_users_email'), table_name='users')
    op.drop_table('users')

Step 2: Review the Migration

Always review the generated migration script to ensure it will make the changes you expect.

Step 3: Apply the Migration

To apply the migration to your database:

bash

alembic upgrade head

This runs all pending migration scripts, updating your database schema to match your models.

Common Migration Operations

Adding a New Column

python

def upgrade():
    op.add_column('users', sa.Column('full_name', sa.String(100), nullable=True))

def downgrade():
    op.drop_column('users', 'full_name')

Renaming a Column

python

def upgrade():
    op.alter_column('users', 'full_name', new_column_name='name')

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

Adding an Index

python

def upgrade():
    op.create_index('ix_users_name', 'users', ['name'])

def downgrade():
    op.drop_index('ix_users_name', 'users')

Adding a Foreign Key

python

def upgrade():
    op.add_column('posts', sa.Column('user_id', sa.Integer(), nullable=True))
    op.create_foreign_key(
        'fk_posts_users', 
        'posts', 'users', 
        ['user_id'], ['id']
    )

def downgrade():
    op.drop_constraint('fk_posts_users', 'posts', type_='foreignkey')
    op.drop_column('posts', 'user_id')

Integrating Alembic with FastAPI

Now let's see how to integrate Alembic with a FastAPI application.

Project Structure

A typical project structure might look like:

my_fastapi_project/
├── alembic/
│   ├── versions/
│   ├── env.py
│   ├── README
│   └── script.py.mako
├── alembic.ini
├── app/
│   ├── __init__.py
│   ├── main.py
│   ├── database.py
│   ├── config.py
│   ├── models.py
│   ├── schemas.py
│   └── crud.py
├── requirements.txt
└── README.md

Database Connection Setup

In app/database.py:

python

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from app.config import settings

SQLALCHEMY_DATABASE_URL = settings.DATABASE_URL

engine = create_engine(SQLALCHEMY_DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

Define Models

In app/models.py:

python

from sqlalchemy import Boolean, Column, Integer, String
from sqlalchemy.orm import relationship

from app.database import Base

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    email = Column(String, unique=True, index=True, nullable=False)
    hashed_password = Column(String, nullable=False)
    is_active = Column(Boolean, default=True)

    # Define relationship
    items = relationship("Item", back_populates="owner")


class Item(Base):
    __tablename__ = "items"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, index=True)
    description = Column(String)
    owner_id = Column(Integer, ForeignKey("users.id"))

    # Define relationship
    owner = relationship("User", back_populates="items")

Create a Migration

After defining your models, create a migration:

bash

alembic revision --autogenerate -m "Create users and items tables"

Apply Migration on Application Startup

You might want to automatically apply migrations when your application starts in some environments. Here's how to do it:

In app/main.py:

python

from fastapi import FastAPI
import subprocess
import sys

from app.database import engine, Base
from app.config import settings

app = FastAPI()

# Apply migrations at startup if configured
if settings.AUTO_APPLY_MIGRATIONS:
    print("Applying database migrations...")
    try:
        subprocess.run(["alembic", "upgrade", "head"], check=True)
        print("Migrations applied successfully!")
    except subprocess.CalledProcessError:
        print("Failed to apply migrations!")
        sys.exit(1)

# Include routers and other app setup
# ...

Real-World Example: Blog Application

Let's create a simple blog application with users, posts, and comments, showing how migrations evolve over time.

Step 1: Initial Model - Users

python

# app/models.py
from sqlalchemy import Column, Integer, String, Boolean
from app.database import Base

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True)
    username = Column(String, unique=True, index=True)
    email = Column(String, unique=True)
    hashed_password = Column(String)
    is_active = Column(Boolean, default=True)

Create the first migration:

bash

alembic revision --autogenerate -m "Create users table"
alembic upgrade head

Step 2: Add Posts

Now, let's add a posts model:

python

# app/models.py (updated)
from sqlalchemy import Column, Integer, String, Boolean, ForeignKey, DateTime
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from app.database import Base

class User(Base):
    # ... (same as before)
    
    # Add relationship
    posts = relationship("Post", back_populates="author")

class Post(Base):
    __tablename__ = "posts"

    id = Column(Integer, primary_key=True)
    title = Column(String, index=True)
    content = Column(String)
    created_at = Column(DateTime, server_default=func.now())
    user_id = Column(Integer, ForeignKey("users.id"))

    # Define relationship
    author = relationship("User", back_populates="posts")

Create a migration for this change:

bash

alembic revision --autogenerate -m "Create posts table"
alembic upgrade head

Step 3: Add Comments

Finally, let's add comments to our blog:

python

# app/models.py (updated again)
# ... (previous code)

class User(Base):
    # ... (same as before)
    
    # Add relationship
    comments = relationship("Comment", back_populates="author")

class Post(Base):
    # ... (same as before)
    
    # Add relationship
    comments = relationship("Comment", back_populates="post")

class Comment(Base):
    __tablename__ = "comments"

    id = Column(Integer, primary_key=True)
    content = Column(String)
    created_at = Column(DateTime, server_default=func.now())
    user_id = Column(Integer, ForeignKey("users.id"))
    post_id = Column(Integer, ForeignKey("posts.id"))

    # Define relationships
    author = relationship("User", back_populates="comments")
    post = relationship("Post", back_populates="comments")

Create a migration for this final change:

bash

alembic revision --autogenerate -m "Create comments table"
alembic upgrade head

Best Practices for Database Migrations

1. Version Control: Always commit migration scripts to your version control system.

2. Test Migrations: Always test migrations on a copy of your production data before deploying.

3. Backup Before Migrating: Always back up your database before running migrations in production.

4. Avoid Direct Schema Changes: Never make manual schema changes. Always use migrations.

5. Organize Migrations: Keep migrations focused and small. One change per migration is often best.

6. Review Generated Migrations: Always review auto-generated migrations before applying them.

7. Downgrade Paths: Ensure downgrade functions work correctly for rollbacks.

8. Data Migrations: When changing schemas, remember to migrate the data too if needed.

Handling Migration Errors

If a migration fails, you can:

1. Fix the error in the migration script
2. Downgrade to the previous version:
   bash

   alembic downgrade -1

3. Apply the fixed migration again:
   bash

   alembic upgrade head

Common Migration Scenarios

1. Adding a Required Column

When adding a required (non-nullable) column to an existing table with data, you need to provide a default value:

python

def upgrade():
    # First add column as nullable
    op.add_column('users', sa.Column('role', sa.String(50), nullable=True))
    
    # Set a default value for existing rows
    op.execute("UPDATE users SET role = 'user' WHERE role IS NULL")
    
    # Then make it non-nullable
    op.alter_column('users', 'role', nullable=False)

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

2. Data Migration

Sometimes you need to transform data during a migration:

python

def upgrade():
    # Add a new column
    op.add_column('users', sa.Column('full_name', sa.String(100), nullable=True))
    
    # Populate it with data from other columns
    op.execute("UPDATE users SET full_name = first_name || ' ' || last_name")
    
    # Drop old columns
    op.drop_column('users', 'first_name')
    op.drop_column('users', 'last_name')

def downgrade():
    # Add back old columns
    op.add_column('users', sa.Column('first_name', sa.String(50), nullable=True))
    op.add_column('users', sa.Column('last_name', sa.String(50), nullable=True))
    
    # Extract data from full_name
    op.execute("UPDATE users SET first_name = split_part(full_name, ' ', 1)")
    op.execute("UPDATE users SET last_name = split_part(full_name, ' ', 2)")
    
    # Drop new column
    op.drop_column('users', 'full_name')

Summary

In this tutorial, you've learned:

1. What database migrations are and why they're important
2. How to set up Alembic for database migrations in FastAPI applications
3. How to create, review, and apply migrations
4. How to handle common migration scenarios
5. Best practices for managing database schema changes

Database migrations are essential for evolving your application's data model while preserving existing data. With Alembic and SQLAlchemy, you have powerful tools to manage these changes in a structured, version-controlled way.

Additional Resources

• Alembic Documentation
• SQLAlchemy Documentation
• FastAPI Documentation
• Database Migrations with Flask-Migrate (Similar principles, different framework)

Exercises

1. Create a FastAPI application with a User model and use Alembic to create the initial migration.
2. Add a new Profile model with a one-to-one relationship to User, and create a migration for it.
3. Modify an existing model by adding a new field and generate a migration.
4. Practice a data migration by transforming data in an existing column.
5. Try to simulate a migration error and practice rolling back to a previous version.

By mastering database migrations, you're well-equipped to build and evolve robust FastAPI applications with confidence.