Database Migrations with sqlmodel and alembic
sqlmodel is a very useful proxy tool that allows pydantic models and sqlalchemy models to be combined. This allows FastAPI, built by the same author, to intuitively know more about the models' metadata.
This guide will show you how to combine this tool into alembic, and create auto-migrations with a set of helper tools.
We will need 2 functions in 2 files, which are pretty simple. The first is the engine creator.
1 2 3 4 5 6 7
# database.py from sqlmodel import create_engine from sqlalchemy.engine import Engine def get_db_engine(url) -> Engine: return create_engine(url)
Which will create the SQLAlchemy Engine object for us. This is used for alembic to connect to the database, I also use this function, cached, to generate Session objects. The next is used for the models.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# model.py import uuid from typing import Optional from sqlmodel import Field, SQLModel class Hero(SQLModel, table=True): id: Optional[int] = Field(default=None, primary_key=True) name: str secret_name: str age: Optional[int] = None def get_metadata(): return SQLModel.metadata
This file contains all our models, registering the metadata as tables, relationships, indexes and anything else that's
get_metadata() method must be after the model definitions so the
SQLModel.metadata is complete.
From here we can touch the alembic env.py function to describe the online environment.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
from myapp.database import get_db_internal from myapp.model import get_metadata ... # `context` defined as before target_metadata = get_metadata() def run_migrations_online() -> None: connectable = get_db_internal() with connectable.connect() as connection: context.configure( connection=connection, target_metadata=target_metadata ) with context.begin_transaction(): context.run_migrations()
The important bit here is the connection and metata are fetched directly from your application via the imports. This will set up all the correct connection details for the environment. To create production migrations you will have to generate the migrations on a host with access the prod environment (i.e. script opening a PR from a prod env).
Now we can continue from normal, running
alembic revision --autogenerate to create revisions on the delta of the
database & metadata Models. Running this will create a new version in the correct directory.
At this point you may need to manually import
sqlmodel to link together column types. Other changes should be made
at this point.
alembic upgrade head will update the database up to the correct revision.