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.
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")
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.
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
downgrade functions are blank - this is where we’ll add the code for our tables.
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
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.
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:
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>.