Files
model-platform/migrations/env.py
T

128 lines
3.4 KiB
Python

from __future__ import annotations
import asyncio
import os
from decimal import Decimal, InvalidOperation
from logging.config import fileConfig
from typing import Any
from alembic import context
from sqlalchemy import Connection, pool
from sqlalchemy.ext.asyncio import async_engine_from_config
from common.db import Base
config = context.config
if config.config_file_name is not None:
fileConfig(config.config_file_name)
target_metadata = Base.metadata
def canonical_default(value: Any) -> tuple[str, Any] | None:
"""Normalize harmless MySQL quoting and numeric formatting differences."""
if value is None:
return None
text_value = str(value).strip()
while (
len(text_value) >= 2
and text_value.startswith("(")
and text_value.endswith(")")
):
text_value = text_value[1:-1].strip()
if (
len(text_value) >= 2
and text_value[0] == text_value[-1]
and text_value[0] in {"'", '"'}
):
text_value = text_value[1:-1]
try:
return ("number", Decimal(text_value).normalize())
except InvalidOperation:
return ("text", text_value.casefold())
def compare_server_default(
migration_context: Any,
inspected_column: Any,
metadata_column: Any,
inspected_default: str | None,
metadata_default: Any,
rendered_metadata_default: str | None,
) -> bool | None:
"""Suppress formatting-only differences and defer real changes to Alembic."""
del migration_context, inspected_column, metadata_column, metadata_default
if canonical_default(inspected_default) == canonical_default(
rendered_metadata_default
):
return False
return None
def database_url() -> str:
"""Return the runtime database URL without storing credentials in the repo."""
try:
return os.environ["DATABASE_URL"]
except KeyError as exc:
raise RuntimeError(
"DATABASE_URL is required for Alembic commands"
) from exc
def configure_context(*, connection: Connection | None = None) -> None:
options = {
"target_metadata": target_metadata,
"compare_type": True,
"compare_server_default": compare_server_default,
}
if connection is None:
context.configure(
url=database_url(),
literal_binds=True,
dialect_opts={"paramstyle": "named"},
**options,
)
else:
context.configure(connection=connection, **options)
def run_migrations_offline() -> None:
configure_context()
with context.begin_transaction():
context.run_migrations()
def do_run_migrations(connection: Connection) -> None:
configure_context(connection=connection)
with context.begin_transaction():
context.run_migrations()
async def run_async_migrations() -> None:
section = config.get_section(config.config_ini_section, {})
section["sqlalchemy.url"] = database_url()
connectable = async_engine_from_config(
section,
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
try:
async with connectable.connect() as connection:
await connection.run_sync(do_run_migrations)
finally:
await connectable.dispose()
def run_migrations_online() -> None:
asyncio.run(run_async_migrations())
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()