Shoin | Let's Share!

FastAPI - SQLAlchemy and Alembic Integration

Abdullah Caliskan -

Almost in every project we store information in databases. It could be an SQL or a NoSQL database. Let's say you would like to go with an SQL database. You have plenty of options to choose. In SQL databases you need to define the structure such as columns and its data types, limitations, default values etc. You may want to create indexes, views and transactions as well. SQLAlchemy allows you to define your tables in form of python classes. You can get the value of a column just by getting the attribute of class. It helps you to get advantage of OOP paradigm. ORM is capable of creating the tables but it's not preferred way for production workloads. It's not a sustainable option. When you modify your structure, it can't make it happen. It's not good as well. Because you don't track what changes you made before. In this point alembic comes to your help. It's capable of creating scripts to modify your structure. It provides rollback scripts as well.

Install required packages

poetry add sqlalchemy alembic
fooapi/ # Main package
    fooapi/
    	__init__.py
        main.py
    tests/
    poetry.lock
    pyproject.toml
    db/
    	__init__.py
        base_class.py
        base.py
        session.py

SQLAlchemy

We should create a new sub-package called db. We store the base classes, session objects and configurations for database connection. The structure is like below.

We define our base class first. We inherit new model classes from the base class. We can define the common columns in base class. It effects every model classes that is inherited from it. So that you don't have to type commun columns in each model classes.

from sqlalchemy.orm import as_declarative
from sqlalchemy import Column, Integer

class_registry: dict = {}

@as_declarative(class_registry=class_registry)
class Base:
	id = Column('id', Integer, primary_key=True, index=True)

class_registry dictionary is used for tracking the which classes are used creating database tables. It's mapping. as_declarative means that we declare the attributes and structure of table. We added an id column which is common for all of our tables.

Now we just defined our Base class. We need to initiate a Session to connect the database.  SQLAlchemy is able to connect to the different database engines with almost same code or with small changes.
We keep database connection string in a variable. We need to create an engine. ORM gets the target engine from connection string. We can pass some extra parameters here. For instance, sqlite doesn't support multi-thread applications. FastAPI creates threads for requests because it's based on starlette (ASGI) which works async. We define that don't check is the connection coming from same thread, or not.
sessionmaker functions returns Session object which is used actually to connect and initiate a session in database. This object is callable. If you want to connect to the db, you call this object. SessionLocal()

from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine

SQL_URI = "sqlite:///app.db"

engine = create_engine(SQL_URI, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(engine, autoflush=False, autocommit=False)

We create another module called base.py I'll explain that in alembic section. This is how it looks like. 

from .base_class import Base

Alembic

We need to initate alembic in project. Go to directory next to pyproject.toml and run alembic init alembic It creates a directory named alembic and alembic.ini file next to pyproject.toml This is last directory setup.

fooapi/ # Main package
    fooapi/
    	__init__.py
        main.py
    tests/
    poetry.lock
    pyproject.toml
    db/
    	__init__.py
        base_class.py
        base.py
        session.py
    # NEW files below
    alembic.ini 
    alembic/
    	versions/
        env.py

Find and comment out sqlalchemy.url line in alembic.ini, because we set this value by getting from session.py Make below changes in env.py

from fooapi.db.base import Base
target_metadata = Base.metadata

from fooapi.db.session import SQL_URI
def get_url() -> str:
	return SQL_URI

target_metadata is required to run alembic properly. It finds Base class get its metadata with table-class mapping that we defined in base_class.py
We imported SQL_URI to point out the target database.

In addition to that, we must change the run_migrations_offline function. We commented out the sqlalchemy.url in alembic.ini. Update the value of url parameter in context.configure method.

context.configure(
    url=get_url,
    target_metadata=target_metadata,
    literal_binds=True,
    dialect_opts={"paramstyle": "named"},
)

We have to make some updates in run_migrations_online function as well. Because we load connection string from session.py We need to set the connection string in this configuration too.

configuration = config.get_section(config.config_ini_section, {})
configuration['sqlalchemy.url'] = get_url()
connectable = engine_from_config(
    configuration,
    prefix="sqlalchemy.",
    poolclass=pool.NullPool,
    )

Alembic Revision and Apply

Alembic creates migration scripts and they're versioned. To create a new revision and apply it, you can give this command

alembic revision --autogenerate -m "first revision"
alembic upgrade head

Most probably it's gonna create anything in database, because we didn't define any model classes that are mapped to real database tables. At least we tested our integration in this step.

Create DB Model

Create a new sub-package called models under the fooapi package. Create a new module named user.py We define the user model class in this module.
It must be inherited from Base class. We have only 2 columns for this table. We have to set the table name as well. 

from fooapi.db.base_class import Base
from sqlalchemy import Column, String


class User(Base):
    __tablename__ = "users"
    
    username = Column(String, index=True)
    password = Column(String)

Now go to the base.py under the db sub-package, import this model class, otherwise alembic isn't going to detect that model class is available.

from .base_class import Base

from fooapi.models.user import User

This part is really important. You must have base.py and base_class.py both.
Alembic needs base.py. Model classes need base_class.py.
Alembic needs Model classes as well, that's why we import model classes in base.py
If you didn't define your Base class in base_class.py you would cause circular dependency. That's why you must define this in separated files. If you get an error about that, ensure you imported Base class from the correct module.

Model Classes;
- Import Base class from base_class.py
- Import Module class within base.py for alembic to discover it.

Alembic;
- Import base.py
- base.py must import Base class from base_class.py
- Make sure you imported all defined Model classes within base.py

To be able to create the table for User model. Give below commands

alembic revision --autogenerate -m "User model added"
alembic upgrade head 


alembic downgrade -1  # To revert changes

Now we have created a table by assist of Alembic. We store the changes in our repository and able to revert and provision versions in database with less-effort.