Skip to content

Instantly share code, notes, and snippets.

@dantetemplar
Last active July 6, 2024 14:59
Show Gist options
  • Save dantetemplar/cbef8b9f1d9d6cde7547629d6d85fcd1 to your computer and use it in GitHub Desktop.
Save dantetemplar/cbef8b9f1d9d6cde7547629d6d85fcd1 to your computer and use it in GitHub Desktop.
Async alembic configuration

Async alembic configuration

Features

Database connection from settings.yaml

You don't need to change alembic.ini to set up a database connection. Alembic will derive database connection uri from settings.yaml (or any yaml which you specify in SETTINGS_PATH environment variable).

Example of settings.yaml:

db_uri: "postgresql+asyncpg://postgres:postgres@localhost:5432/postgres"

Asynchronous driver support

You don't need to separate the database connection setup between your async application and the alembic - just use same async driver.

No empty automigrations

You don't need to manually check and delete empty auto-migrations.

Ruff post-hook which will be applied for each migration

You don't need to run ruff format after each migration, because alembic will do it for you.

Pretty name for migrations

You don't need to remember the migration sequence, now the files are sorted by time and human-readable.

Gist checklist

  1. Clone files from this gist to your project.
  2. Add ruff to your project.
  3. Setup alembic folder:
mkdir -p alembic/versions
mv script.py.mako alembic/
mv env.py alembic/
mv ALEMBIC.md alembic/README.md
touch alembic/versions/.gitkeep
  1. Make sure the sqlalchemy metadata is imported correctly.

Reference

Target database metadata

alembic/env.py

# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
from src.database.models.db import Base  # noqa: E402

target_metadata = Base.metadata

### Read yaml settings
`alembic/env.py` 

```python3
app_settings_path = os.getenv("SETTINGS_PATH", "settings.yaml")
app_settings = safe_load(Path(app_settings_path).read_text())
# get database uri from settings.yaml
config.set_main_option("sqlalchemy.url", app_settings["db_uri"])

Use async driver

alembic/env.py

async def run_async_migrations() -> None:
    """In this scenario we need to create an Engine
    and associate a connection with the context.
    """

    connectable = async_engine_from_config(
        config.get_section(config.config_ini_section, {}),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)

    await connectable.dispose()

Migration file name

alembic.ini

# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
# Uncomment the line below if you want the files to be prepended with date and time
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
# for all available tokens
file_template = _%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d_%%(minute).2d_%%(slug)s

Ruff integration

alembic.ini

# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
hooks = ruff
ruff.type = exec
ruff.executable = poetry
ruff.options = run ruff format REVISION_SCRIPT_FILENAME

No empty automigrations

alembic/env.py

def process_revision_directives(context, revision, directives):
    if config.cmd_opts.autogenerate:
        script = directives[0]
        if script.upgrade_ops.is_empty():
            directives[:] = []
            print("No changes in schema detected.")

License

MIT License

Copyright (c) 2024 Ruslan Belkov

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
# A generic, single database configuration.
[alembic]
# path to migration scripts
script_location = alembic
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
# Uncomment the line below if you want the files to be prepended with date and time
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
# for all available tokens
file_template = _%%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d_%%(minute).2d_%%(slug)s
# sys.path path, will be prepended to sys.path if present.
# defaults to the current working directory.
prepend_sys_path = .
# timezone to use when rendering the date within the migration file
# as well as the filename.
# If specified, requires the python-dateutil library that can be
# installed by adding `alembic[tz]` to the pip requirements
# string value is passed to dateutil.tz.gettz()
# leave blank for localtime
# timezone =
# max length of characters to apply to the
# "slug" field
# truncate_slug_length = 40
# set to 'true' to run the environment during
# the 'revision' command, regardless of autogenerate
# revision_environment = false
# set to 'true' to allow .pyc and .pyo files without
# a source .py file to be detected as revisions in the
# versions/ directory
# sourceless = false
# version location specification; This defaults
# to alembic/versions. When using multiple version
# directories, initial revisions must be specified with --version-path.
# The path separator used here should be the separator specified by "version_path_separator" below.
# version_locations = %(here)s/bar:%(here)s/bat:alembic/versions
# version path separator; As mentioned above, this is the character used to split
# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
# Valid values for version_path_separator are:
#
# version_path_separator = :
# version_path_separator = ;
# version_path_separator = space
version_path_separator = os # Use os.pathsep. Default configuration used for new projects.
# set to 'true' to search source files recursively
# in each "version_locations" directory
# new in Alembic version 1.10
# recursive_version_locations = false
# the output encoding used when revision files
# are written from script.py.mako
# output_encoding = utf-8
[post_write_hooks]
# post_write_hooks defines scripts or Python functions that are run
# on newly generated revision scripts. See the documentation for further
# detail and examples
# format using "black" - use the console_scripts runner, against the "black" entrypoint
# hooks = black
# black.type = console_scripts
# black.entrypoint = black
# black.options = -l 79 REVISION_SCRIPT_FILENAME
# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
hooks = ruff
ruff.type = exec
ruff.executable = poetry
ruff.options = run ruff format REVISION_SCRIPT_FILENAME
# Logging configuration
[loggers]
keys = root,sqlalchemy,alembic
[handlers]
keys = console
[formatters]
keys = generic
[logger_root]
level = WARN
handlers = console
qualname =
[logger_sqlalchemy]
level = WARN
handlers =
qualname = sqlalchemy.engine
[logger_alembic]
level = INFO
handlers =
qualname = alembic
[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic
[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s
datefmt = %H:%M:%S

Alembic

Based on Async alembic configuration with ruff post-hook gist.

Autogenerate migration

alembic revision --autogenerate -m "<some message>"

Upgrade database

Upgrade to the latest revision

alembic upgrade head

Upgrade to the specific revision

alembic upgrade <revision>

Skip migration but apply it in database

Stamp will set the database alembic version to the specified revision but will not change schema.

alembic stamp <revision>
"""
Based on `Async alembic configuration with ruff post-hook` gist:
https://gist.github.com/dantetemplar/cbef8b9f1d9d6cde7547629d6d85fcd1
"""
import asyncio
from logging.config import fileConfig
from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config
from alembic import context
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
fileConfig(config.config_file_name)
# Read settings YAML file
from pathlib import Path # noqa: E402
from yaml import safe_load # noqa: E402
import os # noqa: E402
app_settings_path = os.getenv("SETTINGS_PATH", "settings.yaml")
app_settings = safe_load(Path(app_settings_path).read_text())
# get database uri from settings.yaml
config.set_main_option("sqlalchemy.url", app_settings["db_uri"])
# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
from src.database.models.db import Base # noqa: E402
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.
def run_migrations_offline() -> None:
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
with context.begin_transaction():
context.run_migrations()
def do_run_migrations(connection: Connection) -> None:
context.configure(
connection=connection,
target_metadata=target_metadata,
process_revision_directives=process_revision_directives,
)
with context.begin_transaction():
context.run_migrations()
async def run_async_migrations() -> None:
"""In this scenario we need to create an Engine
and associate a connection with the context.
"""
connectable = async_engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
async with connectable.connect() as connection:
await connection.run_sync(do_run_migrations)
await connectable.dispose()
def run_migrations_online() -> None:
"""Run migrations in 'online' mode."""
connectable = config.attributes.get("connection", None)
if connectable is None:
asyncio.run(run_async_migrations())
else:
do_run_migrations(connectable)
# https://stackoverflow.com/a/71212675/19566814
# almost identical to Flask-Migrate (Thanks miguel!)
# this callback is used to prevent an auto-migration from being generated
# when there are no changes to the schema
def process_revision_directives(context, revision, directives):
if config.cmd_opts.autogenerate:
script = directives[0]
if script.upgrade_ops.is_empty():
directives[:] = []
print("No changes in schema detected.")
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()
"""${message}
Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}
# revision identifiers, used by Alembic.
revision: str = ${repr(up_revision)}
down_revision: Union[str, None] = ${repr(down_revision)}
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
def upgrade() -> None:
${upgrades if upgrades else "pass"}
def downgrade() -> None:
${downgrades if downgrades else "pass"}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment