didier/database/migrations.py

54 lines
1.6 KiB
Python
Raw Normal View History

2022-06-11 01:26:19 +02:00
import logging
2022-07-19 18:49:22 +02:00
from sqlalchemy.orm import Session
2022-07-19 18:57:24 +02:00
from alembic import command, script
2022-07-19 18:49:22 +02:00
from alembic.config import Config
2022-06-11 01:26:19 +02:00
from alembic.runtime import migration
2022-07-25 20:33:20 +02:00
from database.engine import postgres_engine
2022-07-19 18:49:22 +02:00
__config_path__ = "alembic.ini"
__migrations_path__ = "alembic/"
2022-06-11 01:26:19 +02:00
2022-07-19 18:49:22 +02:00
cfg = Config(__config_path__)
cfg.set_main_option("script_location", __migrations_path__)
__all__ = ["ensure_latest_migration", "migrate"]
2022-07-11 22:23:38 +02:00
2022-06-11 01:26:19 +02:00
async def ensure_latest_migration():
"""Make sure we are currently on the latest revision, otherwise raise an exception"""
2022-07-19 18:49:22 +02:00
alembic_script = script.ScriptDirectory.from_config(cfg)
2022-06-11 01:26:19 +02:00
2022-07-25 20:33:20 +02:00
async with postgres_engine.begin() as connection:
2022-06-11 01:26:19 +02:00
current_revision = await connection.run_sync(
lambda sync_connection: migration.MigrationContext.configure(sync_connection).get_current_revision()
)
alembic_head = alembic_script.get_current_head()
if current_revision != alembic_head:
error_message = (
f"Pending migrations (current revision is {current_revision}, while head is at {alembic_head})"
)
logging.error(error_message)
raise RuntimeError(error_message)
2022-07-19 18:49:22 +02:00
def __execute_upgrade(connection: Session):
cfg.attributes["connection"] = connection
command.upgrade(cfg, "head")
def __execute_downgrade(connection: Session):
cfg.attributes["connection"] = connection
command.downgrade(cfg, "base")
async def migrate(up: bool):
"""Migrate the database upwards or downwards"""
2022-07-25 20:33:20 +02:00
async with postgres_engine.begin() as connection:
2022-07-19 18:49:22 +02:00
await connection.run_sync(__execute_upgrade if up else __execute_downgrade)