Code

Python Async Database Migrations With Alembic

This post is about how I currently handle database migrations with Python. My go-to tech stack setup is FastAPI, SQLAlchemy 2.0, and alembic. I also make scripts so that I don’t need to remember the commands, as they are kind of hard to remember sometimes. Those scripts will also be included in this post.

There will be no database setup information, so this post will be assuming that you have a database setup and running already.

If you are only interested in the tech I use, here are the links for them:

Migration Manage Script

Similar to Django’s manage function, I’ve created one for alembic. Since I primarily develop on linux, it is a bash script. But I’m sure it can be converted into a batch script equivalent.

Make a file called migrate.sh:

#!/bin/bash
## Examples:
# ./migrate.sh init
# ./migrate.sh run
# ./migrate.sh add "migration message"
CMD=$1
case $CMD in
    "init")
        alembic init -t async migrations
        ;;
    "run")
        alembic upgrade head
        ;;
    "add")
        if [ -z "$2" ]; then
            echo "Add a name/message for the migration."
            exit 1
        fi
        alembic revision --autogenerate -m "$2"
        ;;
    "revert")
        alembic downgrade -1
        echo "Make sure to delete the downgraded migration file in versions/ folder."
        ;;
    *)
        echo "$1 is not a command.\nCommands: init | run | add {MESSAGE/NAME} | revert"
        ;;
esac

Install Packages

We really only need to install 4 things:

pip install SQLAlchemy alembic python-dotenv

Our 4th package will depend in your database, I’m using postgresql. The model examples below will be using a postgresql specific ID type, make sure to convert it to your DBAPI’s type if you’re not using postgresql:

pip install asyncpg 

In order to find your DBAPI, refer to sqlalchemy’s docs (this links to postgresql) and search for your database of choice and look for it’s available dialects.

Previously, we would need to install the async version, but I believe it is included in SQLAlchemy already.

Database Models

Now we need some models for our database, let’s make something simple like wanting to keep information about our blogs.

# in models/base.py
from sqlalchemy.orm import DeclarativeBase

class Base(DeclarativeBase):
    pass

This Base class is for metadata that lets our migrations know which classes are used for the database.

# in models/blog.py
from models.base import Base
from datetime import datetime
from sqlalchemy import DateTime, String
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.sql import func
from sqlalchemy.orm import Mapped, mapped_column
import uuid

class BlogModel(Base):
    __tablename__ = "blogs"
    id: Mapped[UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    title: Mapped[str] = mapped_column(String(length=160), default="")
    category: Mapped[str] = mapped_column(String(length=60), default="")
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=func.now())
    author: Mapped[str] = mapped_column(String(length=120), default="")
# in models/__init__.py
from models.blog import BlogModel

Setup Alembic

Now we need to setup our migration system for when we make changes to our models. With our manage.sh script we made, run this in the project’s root directory:

./manage.sh migrate init

This will generate something that looks like this:

$ ./manage.sh migrate init
Creating directory '/home/kokopi/01/db-test/migrations' ...  done
Creating directory '/home/kokopi/01/db-test/migrations/versions' ...  done
Generating /home/kokopi/01/db-test/alembic.ini ...  done
Generating /home/kokopi/01/db-test/migrations/README ...  done
Generating /home/kokopi/01/db-test/migrations/env.py ...  done
Generating /home/kokopi/01/db-test/migrations/script.py.mako ...  done
Please edit configuration/connection/logging settings in '/home/kokopi/01/db-test/alembic.ini' before proceeding.

In the alembic.ini:

# uncomment this
file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s

# comment this, we will use environment variables to generate this instead
# sqlalchemy.url = driver://user:pass@localhost/dbname

I’m not sure why these aren’t the default, but having dates on migration files makes sorting way easier. And hard coding the database credentials is always a bad idea.

Create a .env file in the project’s root directory with these values:

DB_HOST="localhost"
DB_NAME="YOUR_DB_NAME"
DB_USERNAME="YOUR_DB_USERNAME"
DB_PASSWORD="YOUR_DB_USERNAME_PASSWORD"

Obviously, replace the caps with your own credentials.

In the migrations/env.py file:

# under this line generated by alembic
config = context.config

# creating our db connection string from environment variables
from dotenv import load_dotenv, dotenv_values
load_dotenv()
envfile = dotenv_values()
env_url = f"postgresql+asyncpg://{envfile['DB_USERNAME']}:{envfile['DB_PASSWORD']}@{envfile['DB_HOST']}/{envfile['DB_NAME']}"
config.set_main_option("sqlalchemy.url", env_url)

# under this line generated by alembic, also can comment it out or delete it
# target_metadata = None

# letting alembic know where our models are
from models.base import Base
from models import *
target_metadata = Base.metadata

Running Migrations

With everything setup, alembic should work. First we need to add a migration file with our script:

$ ./manage.sh migrate add "init"

This should generate an output that looks like this:

$ ./manage.sh migrate add "init"
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.autogenerate.compare] Detected added table 'blogs'
  Generating /home/kokopi/01/db-test/migrations/versions/ddb1890e1613_init.py ... done

If it wasn’t successful, then step through the error to debug it. Maybe it was a typo somewhere or something.

Lastly, we need to run our migration we just made. Using our script:

./manage.sh migrate run

This will generate an output that looks like this:

$ ./manage.sh migrate run
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> ddb1890e1613, init

If everything ran successfully, then congrats. I think there are things that could be improved, but definitely let me know your thoughts once I am able to open up comments here. Or tag me on Bluesky or Twitter.

Since I have recently been creating a few new small projects that uses python and databases, I have created a repo on Github with a similar template if anyone is interested:

Python Async Database Template