Setting up Alembic for SQLAlchemy

🌠 September 5, 2023

Alembic is really useful when you need to manage, version, and migrate your database schema within SQLAlchemy projects. In this post, I want to share with you about how to setup Alembic for a SQLAlchemy project.

First let’s setup the example project:

$ mkdir alembic-example
$ cd alembic-example

Every Python project is suggested to run under virtual environment so we create and activate one.

$ python3 -m venv venv
$ source venv/bin/activate

Next, install the necessary packages. We'll use PostgreSQL in this example, so we'll also install psycopg2:

$ pip3 install sqlalchemy alembic psycopg2

Organize your project with the following structure:

main.py handles basic DB setup, including configuration and connection:

from typing import Dict
from sqlalchemy import create_engine
from sqlalchemy.engine.base import Engine
from sqlalchemy.ext.declarative import declarative_base

ModelBase = declarative_base()
pg_engine: Engine

def setup(config: Dict[str,  str]):
    global pg_engine
    username = config.get('username')
    password = config.get('password')
    database = config.get('database')
    hostname = config.get('hostname')
    port = config.get('port')

    pg_engine = create_engine(
        f'postgresql+psycopg2://{username}:{password}@{hostname}:{port}/{database}',
        echo=True
    )

    return pg_engine

def get_connection():
    global pg_engine
    if pg_engine:
        return pg_engine.connect()
    return None

The models/* directory will contain the model for database interaction.

With our basic setup in place, let's proceed with Alembic initialization:

$ alembic init alembic

This will initializes a new migration environment, which includes the creation of an alembic/ folder. This folder is essential for managing the migration process and contains several critical components:

1. env.py: This script provides runtime instructions to Alembic when it runs migrations. It's where you'd typically specify how Alembic connects to your database and how it should discover your SQLAlchemy models, among other things.
2. script.py.mako: A template for the migration scripts that Alembic generates. Each time you create a new migration, Alembic will use this template to produce a new script inside the versions/ directory.
3. versions/: This is a directory where individual migration scripts are stored. Every time you create a new migration with Alembic (e.g., using alembic revision or alembic revision --autogenerate), a new script will be generated here. Each script has a unique identifier (typically a hash) and contains the logic to upgrade or downgrade your database schema.

Also a **alembic.ini** file will be created at project root which is the main configuration file for Alembic, where you can set various parameters for Alembic's operation, such as the location of the migration script directory, the database URL, and other settings.

After then our project structure looks like this:

├── alembic/
│   ├── README
│   ├── env.py
│   ├── script.py.mako
│   └── versions/
├── alembic.ini
├── main.py
└── models/
    └── __init__.py

Now, let's create the user model at /models/user.py:

# models/user.py
from sqlalchemy import Column, String, Integer
from main import ModelBase

class User(ModelBase):
    __tablename__ = 'user'
    id = Column(Integer, autoincrement=True, primary_key=True)
    name = Column(String(50))
    age = Column(Integer)

And register it in models/__init__.py:

# models/__init__.py
from .user import User

Next we need to do some configuration for Alembic, first open alembic.ini and replace the default value for sqlalchemy.url with your DB URL. I’m using PostgreSQL so I’ll do the following changes:

• Typically, the default Postgres username is your system account, and the default password is an empty string. The default host is localhost, and the default port is 5432.

Update env.py by adding:

from main import ModelBase
from models import *

target_metadata = ModelBase.metadata

Make sure your models are referenced in env.py, otherwise Alembic will have no idea about what your DB schema looks like.

Additionally if you don’t want to hard code sqlalchemy.url, you can do something like this in env.py

DATABASE_URL = os.environ.get('DATABASE_URL')
config.set_main_option('sqlalchemy.url', DATABASE_URL)

With everything set up, execute your first migration:

$ alembic revision --autogenerate -m "Initial migration"
$ alembic upgrade head

Then check your DB and you’ll find that user table is created:

test_db=# \d
                List of relations
 Schema |      Name       |   Type   |   Owner
--------+-----------------+----------+------------
 public | alembic_version | table    | weiguanghe
 public | user            | table    | weiguanghe
 public | user_id_seq     | sequence | weiguanghe
(3 rows)

Let’s add another model role.model:

from sqlalchemy import Column, String, Integer
from main import ModelBase

class Role(ModelBase):
    __tablename__ = 'role'
    id = Column(Integer, autoincrement=True, primary_key=True)
    name = Column(String(50))

Update models/__init__.py accordingly:

from .user import User
from .role import Role

Then, initiate another migration:

$ alembic revision --autogenerate -m "Add role model"
$ alembic upgrade head

Check the database to confirm the addition:

test_db=# \d
                List of relations
 Schema |      Name       |   Type   |   Owner
--------+-----------------+----------+------------
 public | alembic_version | table    | weiguanghe
 public | role            | table    | weiguanghe
 public | role_id_seq     | sequence | weiguanghe
 public | user            | table    | weiguanghe
 public | user_id_seq     | sequence | weiguanghe

It works!

So that’s all about how to setup Alembic for a SQLAlchemy project. Hope this post help you!