derwaro

Async Alembic Notes

Notes on transitioning from a synchronous setup using fastapi, sqlalchemy, and alembic with a sqlite database to a setup that uses fastapi, sqlalchemy, and alembic, but completely async with a mysql db (using aiomysql)

I will only concentrate on the alembic setup and assume a mysql user with sufficient privileges.

  1. install aiomysql

pip install aiomysql

  1. alembic’s env.py
    Assuming a transition from the basic env.py that is generated by running alembic init without any template selected, the file has to be adapted to look something like this. All changes taken from the alembic docs - using asyncio with alembic1 and basically affects “only run_migrations_online will need to be updated”.
 1import asyncio
 2
 3from sqlalchemy.ext.asyncio import async_engine_from_config
 4
 5from logging.config import fileConfig
 6
 7from sqlalchemy import engine_from_config
 8from sqlalchemy import pool
 9
10from alembic import context
11
12# import Base from database
13from src.data.database import Base
14
15# import all Classes from the model files
16from src.models import Class1, Class2
17
18# this is the Alembic Config object, which provides
19# access to the values within the .ini file in use.
20config = context.config
21
22# Interpret the config file for Python logging.
23# This line sets up loggers basically.
24if config.config_file_name is not None:
25    fileConfig(config.config_file_name)
26
27# add your model's MetaData object here
28# for 'autogenerate' support
29# from myapp import mymodel
30# target_metadata = mymodel.Base.metadata
31target_metadata = Base.metadata
32
33# other values from the config, defined by the needs of env.py,
34# can be acquired:
35# my_important_option = config.get_main_option("my_important_option")
36# ... etc.
37
38
39def run_migrations_offline() -> None:
40    """Run migrations in 'offline' mode.
41
42    This configures the context with just a URL
43    and not an Engine, though an Engine is acceptable
44    here as well.  By skipping the Engine creation
45    we don't even need a DBAPI to be available.
46
47    Calls to context.execute() here emit the given string to the
48    script output.
49
50    """
51    url = config.get_main_option("sqlalchemy.url")
52    context.configure(
53        url=url,
54        target_metadata=target_metadata,
55        literal_binds=True,
56        dialect_opts={"paramstyle": "named"},
57    )
58
59    with context.begin_transaction():
60        context.run_migrations()
61
62
63
64def do_run_migrations(connection):
65    context.configure(connection=connection, target_metadata=target_metadata)
66
67    with context.begin_transaction():
68        context.run_migrations()
69
70
71async def run_async_migrations():
72    """In this scenario we need to create an Engine
73    and associate a connection with the context.
74
75    """
76
77    connectable = async_engine_from_config(
78        config.get_section(config.config_ini_section),
79        prefix="sqlalchemy.",
80        poolclass=pool.NullPool,
81    )
82
83    async with connectable.connect() as connection:
84        await connection.run_sync(do_run_migrations)
85
86    await connectable.dispose()
87
88
89def run_migrations_online():
90    """Run migrations in 'online' mode."""
91
92    asyncio.run(run_async_migrations())
93    
94if context.is_offline_mode():
95    run_migrations_offline()
96else:
97    run_migrations_online()
  1. alembic’s alembic.ini
    This value also holds a variable pointing to your database and the driver it should use. Therefore you have to search for a variable called sqlalchemy.url around line 64 and use a connection string that reflects your chosen database type and driver:

sqlalchemy.url = mysql+aiomysql://<db user name>:<db user password>@<hostname>:<port>/<database>

  1. connection string and SQLalchemy engine
    Assuming a set variable SQLALCHEMY_DATABASE_URL in some file like database.py, that references a db and driver that is not the desired mysql/async-driver combination, it is necessary to update this variable to use a mysql db and using the aiomysql driver:

SQLALCHEMY_DATABASE_URL = "mysql+aiomysql://<db user name>:<db user password>@<hostname>:<port>/<database>"

Further, if you used a (sync) sqlite database, it is necessary to update the SQLalchemy engine creation to an async engine (and ommitting the recommended attribute for sqlite databases: connect_args={"check_same_thread": False}):

1from sqlalchemy.ext.asyncio import create_async_engine
2
3engine = create_async_engine(SQLALCHEMY_DATABASE_URL)
  1. update models.py
    If the tables and columns in models.py use some dialect specific instruction, it is necessary to update those to match, in this case, the MYSQL requirements. e.g. sqlite does not enforce length on STRING columns, MYSQL interprets STRING as VARCHAR and wants a maximum length on those columns:

columnname = Column(String(**200**), index=True)
This change needs to be added, I suppose, too to already generated alembic revisions2.

  1. setup environments for alembic
    You probably want to separate environments, speak “development” and “production” or “whatever-you-want”. I’ll leave the database.py part up to you to figure out, since it should be pretty straightforward. For alembic we have to go back to alembic.ini and env.py. There seem to be be various opinions on how to solve this and various ways to solve the distinction between environments.
    One way described in alembic’s cookbook as run multiple alembic environments from one ini file is to have various alembic environments in (one) alembic.ini. This however still implies having database credentials in clear text in this file and having to take care at alembic runtime to specify the desired alembic environment with the --name flag.
    A question asked at stack overflow: is it possible to store the alembic connect string outside of alembic ini provides two possible solutions that don’t include database credentials in clear text and pulls the variables from environment variables. The first proposal seems to leverage an internal alembic API, which is not recommended to do. The second proposal describes interpolation of strings in alembic.ini with variables pulled in env.py from environment. This looks good so far, but does not account for changing environment between “development” and “production”.
    So I clumped together a solution that:

    1. makes changing between “development” and “production” easy
    2. does not store database credentials in neither env.py nor alembic.ini
    3. let’s me run alembic commands without specifying additional flags like --name or --config

    I base my solution on having the url for the development database in alembic.ini’s sqlalchemy.url and only overwrite this value, when running alembic in production mode. All variables are stored in a single place, in my case file config.env.

alembic.ini

1# everything untouched except sqlalchemy.url
2(...)
3# specify the location of your development database, in my case a sqlite db using the async aiosqlite driver
4sqlalchemy.url = sqlite+aiosqlite:///src/sql_app.db
5
6(...)

env.py
based on the already updated version above, some changes for the connection in offline and online are necessary. This is again the complete file, updated and marked with changed/new lines.

  1import asyncio
  2
  3from sqlalchemy.ext.asyncio import async_engine_from_config
  4
  5from logging.config import fileConfig
  6
  7from sqlalchemy import engine_from_config
  8from sqlalchemy import pool
  9
 10from alembic import context
 11
 12# import Base from database
 13from src.data.database import Base
 14
 15# import all Classes from the model files
 16from src.models import Class1, Class2
 17
 18# import the config.env file which holds variables for the environment and the database url parts
 19from src.config import get_config
 20
 21# import the actual values from config.env
 22envconfig = get_config()
 23
 24# this is the Alembic Config object, which provides
 25# access to the values within the .ini file in use.
 26config = context.config
 27
 28# Interpret the config file for Python logging.
 29# This line sets up loggers basically.
 30if config.config_file_name is not None:
 31    fileConfig(config.config_file_name)
 32
 33# add your model's MetaData object here
 34# for 'autogenerate' support
 35# from myapp import mymodel
 36# target_metadata = mymodel.Base.metadata
 37target_metadata = Base.metadata
 38
 39# other values from the config, defined by the needs of env.py,
 40# can be acquired:
 41# my_important_option = config.get_main_option("my_important_option")
 42# ... etc.
 43
 44
 45def run_migrations_offline() -> None:
 46    """Run migrations in 'offline' mode.
 47
 48    This configures the context with just a URL
 49    and not an Engine, though an Engine is acceptable
 50    here as well.  By skipping the Engine creation
 51    we don't even need a DBAPI to be available.
 52
 53    Calls to context.execute() here emit the given string to the
 54    script output.
 55
 56    """
 57    if envconfig.ENVIRONMENT == "development":
 58        url = config.get_main_option("sqlalchemy.url")
 59    elif envconfig.ENVIRONMENT == "production":
 60        url = f"mysql+aiomysql://{envconfig.DB_USER}:{envconfig.DB_PASSWORD}@{envconfig.DB_HOST}:{envconfig.DB_PORT}/{envconfig.DB_DATABASE}"
 61
 62    context.configure(
 63        url=url,
 64        target_metadata=target_metadata,
 65        literal_binds=True,
 66        dialect_opts={"paramstyle": "named"},
 67    )
 68
 69    with context.begin_transaction():
 70        context.run_migrations()
 71
 72
 73def do_run_migrations(connection):
 74    context.configure(connection=connection, target_metadata=target_metadata)
 75
 76    with context.begin_transaction():
 77        context.run_migrations()
 78
 79
 80async def run_async_migrations():
 81    """In this scenario we need to create an Engine
 82    and associate a connection with the context.
 83
 84    """
 85    if envconfig.ENVIRONMENT == "development":
 86        connectable = async_engine_from_config(
 87            config.get_section(config.config_ini_section),
 88            prefix="sqlalchemy.",
 89            poolclass=pool.NullPool,
 90        )
 91    elif envconfig.ENVIRONMENT == "production":
 92        async_config = {
 93            "sqlalchemy.url": f"mysql+aiomysql://{envconfig.DB_USER}:{envconfig.DB_PASSWORD}@{envconfig.DB_HOST}:{envconfig.DB_PORT}/{envconfig.DB_DATABASE}",
 94            "sqlalchemy.poolclass": "NullPool",
 95        }
 96        connectable = async_engine_from_config(
 97            async_config,
 98            prefix="sqlalchemy.",
 99            poolclass=pool.NullPool,
100        )
101
102    async with connectable.connect() as connection:
103        await connection.run_sync(do_run_migrations)
104
105    await connectable.dispose()
106
107
108def run_migrations_online():
109    """Run migrations in 'online' mode."""
110
111    asyncio.run(run_async_migrations())
112
113
114if context.is_offline_mode():
115    run_migrations_offline()
116else:
117    run_migrations_online()

Some explanations:
Lines 18-22: I use a function to import my config from the file config.env. These lines here simply import the function that returns my environment variables, which I can, as can be seen further down, be used like envconfig.<VARIABLENAME>.
Lines 57-60: These lines check, the ENVIRONMENT variable set in config.env. In “development” alembic falls back to the variable set in alembic.ini at sqlalchemy.url3. If I set ENVIRONMENT in config.env to “production”, I use an f-string to create the full connection string (all secrets/variables pulled from config.env)
Lines 85-100: The section for “development” is pretty straightforward to set up, since it, again, falls back to alembic.ini’s sqlalchemy.url. The section for “production” was a bit trickier. I solved it by defining a dictionary, mimicking the structure in alembic.ini, that holds the same connection string as in the offline section. This dictionary is passed to async_engine_from_config, which extracts the safe connection string form it and creates an sqlalchemy async engine, which in turn is used by alembic.

  1. That’s all folks

I hope I captured all my modifications.

  1. Alembic’s cookbook is delicious by the way! ↩︎

  2. to be honest, I simply scrapped all existing revisions and create a new one, that included those changes. ↩︎

  3. Since I use a sqlite database for development, there is no risk of exposing database credentials. For other databases it would be necessary to infer the credentials like on line 60 and completly ignore/not set sqlalchemy.url. ↩︎

Tags: