Basics of Database Migrations with SQLAlchemy and Alembic

PUBLISHED ON DEC 2, 2020 — ALEMBIC, PYTHON, SQLALCHEMY

In this post I’m going to cover the basics of performing database migrations for your Python apps using SQLAlchemy and Alembic.
This is something I did on my stream and thought might be useful to others. The code is available on my GitHub repo.

I’m starting with an existing API written in FastAPI, so won’t be covering setting that up in this post.

First things first, we’ll take a look at the models set up in the API - since I’m keeping this simple, there are just two models.
There’s the Company model:

from sqlalchemy import Integer, String
from sqlalchemy.sql.schema import Column
from sqlalchemy.orm import relationship
from ..database import Base


class Company(Base):
    __tablename__ = 'companies'

    id = Column(Integer, primary_key=True)
    user_id = Column(String, nullable=False)
    name = Column(String, nullable=False)

    jobs = relationship("Job", back_populates="company")

And the Job model:

from sqlalchemy import Integer, String, ForeignKey
from sqlalchemy.sql.schema import Column
from sqlalchemy.orm import relationship
from ..database import Base


class Job(Base):
    __tablename__ = 'jobs'

    id = Column(Integer, primary_key=True)
    title = Column(String, nullable=False)
    description = Column(String, nullable=False)

    company_id = Column(Integer, ForeignKey('companies.id'))

There’s a one-to-many relationship set up there with SQLAlchemy - one company can have many jobs.
At this point, we have the models set up, but we don’t have anything in the database - though we have a connection set up.
I’ve added a Docker Compose file to run the API and Database to hopefully make things easier, so you can run docker-compose up in the root of the directory to spin everything up. Running the compose file will create the database we’re using for this example.

So now we want to actually create the tables for those models, and to do that we’ll need alembic, so install that with pip pip install alembic, or if you’re installing all the dependencies for the project at once use the requirements.txt file by running pip install -r requirements.txt.

Now we have alembic we need to set up the environment. Make sure you’re in the root of the API project and run the command alembic init alembic.
This creates the alembic environment in our project - you’ll see an alembic directory where the migrations will be stored, and the env config, and in the root of the API directory you’ll see the alembic.ini config file.
The alembic.ini file needs to be updated for the connection string to our database, so look for the sqlalchemy.url property to point it to the database running on Docker: sqlalchemy.url = postgresql://postgres:mysuperpassword@localhost/myawesomedb

And then we can create our first migration by running alembic revision -m "init" - change the name to something else if you prefer.

If you open the migration file, you can see that the upgrade and downgrade functions are blank - this is where we’ll add the code for our tables. In upgrade we want to create our tables, starting with the companies table. We need to call create_table and pass in the name of the table, companies for this example, and the definition of each column - since this is a fairly simple example, I’m adding a primary key id column - this will auto increment by default - and a name column:

    op.create_table(
        'companies',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('name', sa.String(50), nullable=False),
    )

Then we need to create the jobs table, and it’s in here that we need to specify our foreign key to set up the relationship.
As above, we create the table passing in the name and column definitions. The main thing we want to note here is that we’re defining a company_id column which is going to be our foreign key.
Then we have to add the foreign key constraint for that column using ForeignKeyConstraint - the first parameter is the column we’re applying the constraint to, and the second is the relation we want to make in the form of <TABLE>.<COLUMN>, which in our case is companies.id.

op.create_table(
        'jobs',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('description', sa.String(), nullable=False),
        sa.Column('title', sa.String(100), nullable=False),
        sa.Column('company_id', sa.Integer, nullable=False),
        sa.ForeignKeyConstraint(('company_id',), ['companies.id'], ),
    )

This defines how to create our two tables, and how to set up the one-to-many relationship between them.

The downgrade function is simpler as we just want to drop the tables, we just drop them in reverse order to how they were created because of the relationship:

    op.drop_table('jobs')
    op.drop_table('companies')

Now we’ve defined our migration, we need to run them against the database - so make sure your database is up and running.
Make sure you’re in the root of the API directory and run the command: alembic upgrade head

If you connect to the database you’ll now be able to see the two created tables, and the relationship between them.

If you want to undo the migration, you can run the command: alembic downgrade -1
This will undo the last migration applied.
If instead you have multiple migrations and you want to undo a specific one, you need to find the id of the migration by running alembic history. Grab the id, and then use the command alembic downgrade <MIGRATION_ID>.

Hopefully this post was useful to someone.
If you have any questions feel free to reach out to me on Twitter or drop into the Twitch stream.