Using SQLModel with Alembic

TL;DR

1. Install Alembic and SQLModel.

2. Create models with SQLModel.SQLModel base class.

3. Add naming conventions to SQLModel.metadata using SQLModel.metadata.naming_convention.

4. Initialize Alembic with alembic init migrations.

5. Add your database URL to the alembic.ini configuration file.

6. In migrations/env.py, set target_metadata to SQLModel.metadata. This should be done after the Model classes have been imported.

7. Add the render_as_batch=True parameter to context.configure(...) in the function run_migrations_online of the migrations/env.py file.

8. Add the user_module_prefix="sqlmodel.sql.sqltypes." parameter to context.configure(...) in the function run_migrations_online of the migrations/env.py file.

9. Add import sqlmodel.sql.sqltypes in migrations/script.py.mako.

10. Now make changes in your data models and run alembic revision --autogenerate -m "some change message".

11. Upgrade the database using alembic upgrade head.

The Goal

We want to create a basic Pet manager app.

Setting Up the Environment

This step is optional.

We use Poetry for package management.

Our project environment looks like this:

sqlmodel-alembic/
├── sqlmodel_alembic/
│   ├── __init__.py
│   └── models.py
├── poetry.lock
└── pyproject.toml

At the time of writing, the dependencies are:

# pyproject.toml
[tool.poetry.dependencies]
python = "^3.8"
sqlmodel = "^0.0.8"
alembic = "^1.9.2"

Installing Alembic and SQLModel

Install alembic and SQLModel with your preferred package manager. I prefer Poetry.

poetry add alembic SQLModel

Make sure you activate the Poetry shell using poetry shell before executing any Alembic command.

Creating Models Using SQLModel

We create two entities: User and Pet.

# sqlmodel_alembic/models.py
from typing import Optional
from sqlmodel import Field, SQLModel

metadata = SQLModel.metadata


class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str


class Pet(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str

SQL Constraint Naming Convention

This step is optional but recommended.

Consider utilising a naming convention when adding SQL constraints to help prevent migration errors (especially with SQLite).

# sqlmodel_alembic/models.py
# ...
NAMING_CONVENTION = {
    "ix": "ix_%(column_0_label)s",
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s",
}

metadata = SQLModel.metadata
metadata.naming_convention = NAMING_CONVENTION
# ...

Alembic's documentation explains the importance of naming conventions.

Initialize Alembic Migrations Directory

alembic init migrations
# Creating directory ./sqlmodel-alembic/migrations ...  done
# Creating directory ./sqlmodel-alembic/migrations/versions ...  done
# Generating ./sqlmodel-alembic/migrations/README ...  done
# File ./sqlmodel-alembic/alembic.ini already exists, skipping
# Generating ./sqlmodel-alembic/migrations/env.py ...  done
# Generating ./sqlmodel-alembic/migrations/script.py.mako ...  done
#Please edit configuration/connection/logging settings in './sqlmodel-alembic/alembic.ini' before proceeding.

This generates alembic.ini and migrations directory. After executing the command, our project directory looks like this:

sqlmodel-alembic/
├── migrations/
│   ├── versions/
│   ├── env.py
│   ├── README
│   └── script.py.mako
├── sqlmodel_alembic/
│   ├── __init__.py
│   └── models.py
├── alembic.ini
├── poetry.lock
└── pyproject.toml

Add Database URL to alembic.ini

# alembic.ini
[alembic]
# ... some other configuration above 👆
sqlalchemy.url = sqlite:///./data.db
# ... some other configuration below 👇

A better way to do this is by getting the database URL from the environment. However, that is out of the scope of this tutorial.

Set target metadata

# migrations/env.py
# ...
import sqlmodel_alembic.models
target_metadata = sqlmodel_alembic.models.metadata
# ...

Note that we use the metadata variable from the sqlmodel_alembic.models module itself. Why? We want metadata to store information about our models during runtime.

Configure env.py

Batch Rendering Mode

Miguel Grinberg's blog outlines the importance of batch rendering and why it is NOT a magic bullet.

With SQLite, the ALTER TABLE command only supports adding or renaming columns. Any other change to columns or constraints is going to be rejected with an error.

We instruct Alembic to perform batch operations by configuring run_migrations_online() like this:

# migrations/env.py
# ...
def run_migrations_online() -> None:
    # ...
    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
+           render_as_batch=True,
        )
# ...

You can also include render_as_batch=True in run_migrations_offline, however, this has not been tested.

What is the difference between offline and online migration?

Taken from Alembic's documentation:

Generation of migrations as SQL scripts, instead of running them against the database - is also referred to as offline mode.

SQLModel Custom Datatypes

SQLModel uses its custom datatypes and we need to instruct Alembic to include them in migration scripts.

For that, we add the use_module_prefix argument to context.configure() in run_migrations_online().

# ...
def run_migrations_online() -> None:
    # ...
    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            render_as_batch=True,
+           user_module_prefix="sqlmodel.sql.sqltypes.",
        )
# ...

You can also include user_module_prefix="sqlmodel.sql.sqltypes." in run_migrations_offline, however, this has not been tested.

For more information, refer to Alembic's documentation on custom datatypes.

Adding SQLModel Import to script.py.mako

Finally, we include the actual SQLModel datatypes import in migrations/script.py.mako.

# ...
from alembic import op
import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
${imports if imports else ""}
# ...

For more information, refer to Alembic's documentation on custom datatypes.

Create the Initial Migration

alembic revision --autogenerate -m "initial"
# INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
# INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
# INFO  [alembic.autogenerate.compare] Detected added table 'pet'
# INFO  [alembic.autogenerate.compare] Detected added table 'user'
#  Generating ./sqlmodel-alembic/migrations/versions/<somehash>_initial.py ...  done

The generated migration script should look approximately like this:

# migrations/versions/<somehash>_initial.py
"""initial

Revision ID: <somehash>
Revises: 
Create Date: 2023-01-29 19:39:24.556269

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes


# revision identifiers, used by Alembic.
revision = '<somehash>'
down_revision = None
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('pet',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_table('user',
    sa.Column('id', sa.Integer(), nullable=False),
    sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('user')
    op.drop_table('pet')
    # ### end Alembic commands ###

Migrate our Database

alembic upgrade head
# INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
# INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
# INFO  [alembic.runtime.migration] Running upgrade  -> <somehash>, initial

And done! You are now harnessing the power of Alembic and SQLModel to manage your database.

What's next?

Next, we are going to use SQLModel with Asyncio.