Using SQLModel with Alembic
Tags: python, sql, alembic, fastapi
Published: 2023-01-29 20:36:35

## 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](https://python-poetry.org/) for package management.

Our project environment looks like this:

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

At the time of writing, the dependencies are:

```ini
# 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](https://python-poetry.org/).

```bash
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.

```python
# 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).

```python
# 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](https://alembic.sqlalchemy.org/en/latest/naming.html) explains the importance of naming conventions.

## Initialize Alembic Migrations Directory

```bash
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:

```plaintext
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`

```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

```python
# 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](https://blog.miguelgrinberg.com/post/fixing-alter-table-errors-with-flask-migrate-and-sqlite) 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:

```diff
# 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](https://alembic.sqlalchemy.org/en/latest/offline.html):

> 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()`.

```diff
# ...
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.](https://alembic.sqlalchemy.org/en/latest/autogenerate.html#comparing-and-rendering-types)

## Adding SQLModel Import to `script.py.mako`

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

```diff
# ...
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.](https://alembic.sqlalchemy.org/en/latest/autogenerate.html#comparing-and-rendering-types)

## Create the Initial Migration

```bash
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:

```python
# 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

```bash
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.