550 lines
20 KiB
Python
550 lines
20 KiB
Python
"""Alembic wiring: scaffolding, autogenerate and applying migrations.
|
|
|
|
`backend/alembic.ini` and `backend/migrations/` are generated on first use and are
|
|
never overwritten, so the plain `alembic` CLI works alongside `db_setup.init_db()`.
|
|
Models are discovered automatically: every `<package>/models.py` under `backend/`
|
|
is imported before the metadata is diffed against the live schema.
|
|
|
|
python alembic_setup.py [migrate|revision|makemigrations|upgrade|downgrade|current|head|stamp] [-m MSG] [-r REV]
|
|
|
|
Django-shaped aliases (same behaviour, different names):
|
|
|
|
python alembic_setup.py makemigrations -m "add form_data"
|
|
python alembic_setup.py upgrade # apply versions/*.py (like migrate)
|
|
python alembic_setup.py stamp -r f3a7e5b34c86 # bookmark only; no DDL
|
|
|
|
Named `alembic_setup` rather than `alembic` because `backend/` is on `sys.path`,
|
|
so a module called `alembic.py` would shadow the installed package.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import logging
|
|
import sys
|
|
from contextlib import asynccontextmanager
|
|
from importlib import import_module
|
|
from pathlib import Path
|
|
from typing import Any, AsyncIterator, Callable, Sequence
|
|
|
|
from alembic import command
|
|
from alembic.autogenerate import compare_metadata, produce_migrations, render_python_code
|
|
from alembic.config import Config
|
|
from alembic.operations import Operations
|
|
from alembic.runtime.migration import MigrationContext
|
|
from alembic.script import ScriptDirectory
|
|
from alembic.script.revision import ResolutionError
|
|
from alembic.util import rev_id as new_rev_id
|
|
from alembic.util.exc import CommandError
|
|
from sqlalchemy import MetaData, text
|
|
from sqlalchemy.engine import Connection
|
|
|
|
from db_setup import BASE_DIR, Base, close_db, database_url, get_engine, get_settings
|
|
|
|
logger = logging.getLogger("db.alembic")
|
|
|
|
INI = BASE_DIR / "alembic.ini"
|
|
MIGRATIONS = BASE_DIR / "migrations"
|
|
VERSIONS = MIGRATIONS / "versions"
|
|
SKIP_DIRS = {"migrations", "__pycache__", "tests", "test", ".venv", "venv", "node_modules"}
|
|
LOCK_ID = 8_412_557_390_112_004 # any 64-bit constant; serialises startup migrations
|
|
|
|
INI_TEMPLATE = """\
|
|
# Generated by alembic_setup.py. The URL is injected from the environment at runtime.
|
|
[alembic]
|
|
script_location = migrations
|
|
prepend_sys_path = .
|
|
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
|
timezone = UTC
|
|
"""
|
|
|
|
ENV_TEMPLATE = '''"""Alembic environment -- generated by alembic_setup.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from alembic import context
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
import alembic_setup as setup # noqa: E402
|
|
import db_setup # noqa: E402
|
|
|
|
metadata = setup.target_metadata()
|
|
options = setup.context_options()
|
|
|
|
|
|
def run(connection) -> None:
|
|
context.configure(connection=connection, target_metadata=metadata, **options)
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
|
|
|
|
if context.is_offline_mode():
|
|
context.configure(
|
|
url=db_setup.database_url(async_driver=False),
|
|
target_metadata=metadata,
|
|
literal_binds=True,
|
|
dialect_opts={"paramstyle": "named"},
|
|
**options,
|
|
)
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
elif (connection := context.config.attributes.get("connection")) is not None:
|
|
run(connection) # alembic_setup passed an already-open connection
|
|
else:
|
|
asyncio.run(setup.run_standalone(run)) # bare `alembic` CLI
|
|
'''
|
|
|
|
MAKO_TEMPLATE = '''"""${message}
|
|
|
|
Revision ID: ${up_revision}
|
|
Revises: ${down_revision | comma,n}
|
|
Create Date: ${create_date}
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
import sqlmodel # SQLModel renders AutoString() into migrations but adds no import
|
|
${imports if imports else ""}
|
|
|
|
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"}
|
|
'''
|
|
|
|
|
|
def scaffold() -> None:
|
|
"""Create the Alembic layout if missing; existing files are left alone."""
|
|
VERSIONS.mkdir(parents=True, exist_ok=True)
|
|
(VERSIONS / ".gitkeep").touch()
|
|
for path, content in (
|
|
(INI, INI_TEMPLATE),
|
|
(MIGRATIONS / "env.py", ENV_TEMPLATE),
|
|
(MIGRATIONS / "script.py.mako", MAKO_TEMPLATE),
|
|
):
|
|
if not path.exists():
|
|
path.write_text(content, encoding="utf-8")
|
|
logger.info("created %s", path)
|
|
|
|
|
|
_imported = False
|
|
|
|
|
|
def target_metadata() -> MetaData:
|
|
"""`Base.metadata` with every discovered model module imported onto it."""
|
|
global _imported
|
|
if _imported:
|
|
return Base.metadata
|
|
if str(BASE_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BASE_DIR))
|
|
names = get_settings().db_model_modules or [
|
|
".".join(p.relative_to(BASE_DIR).with_suffix("").parts).removesuffix(".__init__")
|
|
for p in [*BASE_DIR.rglob("models.py"), *BASE_DIR.rglob("models/__init__.py")]
|
|
if not SKIP_DIRS & set(p.relative_to(BASE_DIR).parts)
|
|
]
|
|
for name in names:
|
|
try:
|
|
import_module(name)
|
|
except Exception as exc:
|
|
logger.warning("skipping model module %s: %s", name, exc)
|
|
_imported = True
|
|
return Base.metadata
|
|
|
|
|
|
def config(connection: Connection | None = None) -> Config:
|
|
scaffold()
|
|
cfg = Config(str(INI))
|
|
cfg.set_main_option("script_location", str(MIGRATIONS))
|
|
# ConfigParser interpolates '%', which is legal inside a password.
|
|
cfg.set_main_option("sqlalchemy.url", database_url(async_driver=False).replace("%", "%%"))
|
|
if connection is not None:
|
|
cfg.attributes["connection"] = connection
|
|
return cfg
|
|
|
|
|
|
VERSION_TABLE = "alembic_version"
|
|
MANUAL_TABLE = "manual_migrations"
|
|
|
|
|
|
def _include_object(obj: Any, name: str, type_: str, reflected: bool, compare_to: Any) -> bool:
|
|
"""Keep autogenerate inside the schemas this application owns."""
|
|
s = get_settings()
|
|
if type_ == "foreign_key_constraint":
|
|
# Reflected FKs are unqualified (search_path=app) while models use
|
|
# app.table; Alembic reports every FK as drop+add. Real FK changes
|
|
# go through models + migrate/manual SQL, not autogenerate.
|
|
return False
|
|
if type_ != "table":
|
|
return True
|
|
if name in (VERSION_TABLE, MANUAL_TABLE): # migration bookkeeping; never ours to alter
|
|
return False
|
|
return not s.db_schemas or (obj.schema or s.db_default_schema) in s.db_schemas
|
|
|
|
|
|
def _skip_empty(context_: Any, revision: Any, directives: list[Any]) -> None:
|
|
"""Write no revision file at all when the models and the database already agree."""
|
|
ops = directives[0].upgrade_ops if directives else None
|
|
if ops is not None and ops.is_empty():
|
|
directives[:] = []
|
|
logger.info("no model changes detected")
|
|
|
|
|
|
def context_options() -> dict[str, Any]:
|
|
"""Shared by `migrations/env.py` and the in-process drift check."""
|
|
s = get_settings()
|
|
return {
|
|
"compare_type": True,
|
|
# Server-default string forms differ between reflection and models
|
|
# (e.g. now() vs CURRENT_TIMESTAMP); comparing them re-applies the same
|
|
# ALTER on every boot when we sync filelessly.
|
|
"compare_server_default": False,
|
|
"include_schemas": bool(s.db_schemas),
|
|
"version_table_schema": s.db_default_schema or None,
|
|
"include_object": _include_object,
|
|
"process_revision_directives": _skip_empty,
|
|
}
|
|
|
|
|
|
async def _run(fn: Callable[[Connection], Any]) -> Any:
|
|
"""Run a synchronous Alembic call on the async engine's connection."""
|
|
schema = get_settings().db_default_schema or "public"
|
|
async with get_engine().connect() as conn:
|
|
# Unqualified FK targets (REFERENCES users) must resolve in `app`.
|
|
await conn.execute(text(f'SET search_path TO "{schema}", public'))
|
|
result = await conn.run_sync(fn)
|
|
await conn.commit()
|
|
return result
|
|
|
|
|
|
async def run_standalone(fn: Callable[[Connection], Any]) -> None:
|
|
"""Entry point for the bare `alembic` CLI, which brings no connection of its own."""
|
|
try:
|
|
await _run(fn)
|
|
finally:
|
|
await close_db()
|
|
|
|
|
|
def head() -> str | None:
|
|
"""The latest revision on disk."""
|
|
heads = ScriptDirectory.from_config(config()).get_heads()
|
|
return heads[0] if heads else None
|
|
|
|
|
|
async def current() -> str | None:
|
|
"""The revision the database is stamped with."""
|
|
opts = {"version_table_schema": get_settings().db_default_schema or None}
|
|
return await _run(lambda c: MigrationContext.configure(c, opts=opts).get_current_revision())
|
|
|
|
|
|
MODELS_STAMP = "models" # alembic_version marker when no revision files ship in the image
|
|
|
|
|
|
def _revision_on_disk(revision_id: str) -> bool:
|
|
"""True when `revision_id` exists under migrations/versions (or is a known alias)."""
|
|
try:
|
|
ScriptDirectory.from_config(config()).get_revision(revision_id)
|
|
except (ResolutionError, CommandError):
|
|
# ScriptDirectory.get_revision wraps ResolutionError in CommandError.
|
|
return False
|
|
return True
|
|
|
|
|
|
async def upgrade(revision: str = "head") -> None:
|
|
if revision == "head" and not head():
|
|
logger.info("no alembic revisions on disk; skipping upgrade")
|
|
return
|
|
# versions/*.py are gitignored, so each machine (and RDS) can stamp an id
|
|
# this checkout has never seen. Alembic then dies with ResolutionError
|
|
# before any DDL. Skip rather than crash; apply_model_drift still runs
|
|
# when DB_AUTOGENERATE is on.
|
|
current_rev = await current()
|
|
if current_rev and not _revision_on_disk(current_rev):
|
|
logger.warning(
|
|
"database revision %s is not in migrations/versions/; skipping alembic upgrade",
|
|
current_rev,
|
|
)
|
|
return
|
|
await _run(lambda c: command.upgrade(config(c), revision))
|
|
logger.info("upgraded to %s", revision)
|
|
|
|
|
|
async def stamp(revision: str = "head") -> None:
|
|
target = revision
|
|
if target == "head" and not head():
|
|
target = MODELS_STAMP
|
|
await _run(lambda c: command.stamp(config(c), target))
|
|
logger.info("stamped database at %s", target)
|
|
|
|
|
|
async def _schema_is_empty() -> bool:
|
|
"""True when the app schema has never been populated (fresh Compose volume)."""
|
|
schema = get_settings().db_default_schema or "public"
|
|
async with get_engine().connect() as conn:
|
|
row = (
|
|
await conn.execute(
|
|
text(
|
|
"SELECT 1 FROM information_schema.tables "
|
|
"WHERE table_schema = :schema AND table_name = 'users' LIMIT 1"
|
|
),
|
|
{"schema": schema},
|
|
)
|
|
).first()
|
|
return row is None
|
|
|
|
|
|
async def bootstrap_empty() -> None:
|
|
"""Create every model table and stamp a revision marker.
|
|
|
|
Several historical revisions assume tables (e.g. job_posts) that were never
|
|
given a create_table in the chain — they only exist on DBs that grew via
|
|
autogenerate. A brand-new Compose Postgres volume therefore cannot
|
|
`upgrade head`. Creating from metadata then stamping is the production
|
|
bootstrap for that case; existing databases keep the normal upgrade path.
|
|
|
|
Revision `.py` files are gitignored and excluded from images; stamp uses
|
|
`models` when the versions directory is empty.
|
|
"""
|
|
metadata = target_metadata()
|
|
async with get_engine().begin() as conn:
|
|
await conn.run_sync(metadata.create_all)
|
|
await stamp("head")
|
|
logger.info("bootstrapped empty database from models")
|
|
|
|
|
|
async def downgrade(revision: str = "-1") -> None:
|
|
await _run(lambda c: command.downgrade(config(c), revision))
|
|
logger.info("downgraded to %s", revision)
|
|
|
|
|
|
def _apply_upgrade_ops(connection: Connection) -> int:
|
|
"""Apply ORM→DB diffs in-process without writing a revision file."""
|
|
opts = {k: v for k, v in context_options().items() if k != "process_revision_directives"}
|
|
ctx = MigrationContext.configure(connection, opts=opts)
|
|
script = produce_migrations(ctx, target_metadata())
|
|
if script.upgrade_ops.is_empty():
|
|
return 0
|
|
operations = Operations(ctx)
|
|
applied = 0
|
|
stack = [script.upgrade_ops]
|
|
while stack:
|
|
elem = stack.pop(0)
|
|
if hasattr(elem, "ops"):
|
|
stack.extend(elem.ops)
|
|
else:
|
|
operations.invoke(elem)
|
|
applied += 1
|
|
return applied
|
|
|
|
|
|
async def apply_model_drift() -> bool:
|
|
"""Sync the live schema to the ORM without creating migration files.
|
|
|
|
Used by Docker/prod boots so `versions/*.py` can stay gitignored and out of
|
|
the image. Returns True when at least one DDL op was applied.
|
|
"""
|
|
opts = {k: v for k, v in context_options().items() if k != "process_revision_directives"}
|
|
diffs = await _run(
|
|
lambda c: compare_metadata(MigrationContext.configure(c, opts=opts), target_metadata())
|
|
)
|
|
if not diffs:
|
|
logger.info("schema matches the models")
|
|
return False
|
|
logger.info("%s schema difference(s) detected; applying without revision files", len(diffs))
|
|
applied = await _run(_apply_upgrade_ops)
|
|
logger.info("applied %s schema operation(s)", applied)
|
|
return applied > 0
|
|
|
|
|
|
def _write_revision(connection: Connection, message: str) -> str | None:
|
|
"""Write a versions/*.py file from the current ORM→DB diff.
|
|
|
|
Uses the on-disk head as down_revision and never reads alembic_version, so
|
|
a stamp from another machine (versions are gitignored) cannot block
|
|
makemigrations. The live database bookmark is left unchanged.
|
|
"""
|
|
opts = {k: v for k, v in context_options().items() if k != "process_revision_directives"}
|
|
ctx = MigrationContext.configure(connection, opts=opts)
|
|
script = produce_migrations(ctx, target_metadata())
|
|
if script.upgrade_ops.is_empty():
|
|
return None
|
|
script_dir = ScriptDirectory.from_config(config(connection))
|
|
revid = new_rev_id()
|
|
script_dir.generate_revision(
|
|
revid,
|
|
message,
|
|
head=script_dir.get_current_head() or "base",
|
|
upgrades=render_python_code(script.upgrade_ops, migration_context=ctx),
|
|
downgrades=render_python_code(script.downgrade_ops, migration_context=ctx),
|
|
)
|
|
return revid
|
|
|
|
|
|
async def autogenerate(message: str = "auto") -> str | None:
|
|
"""Write a revision if the models have drifted; return its id, or None.
|
|
|
|
Local/CLI only. Docker boots use `apply_model_drift` instead so revision
|
|
files are never written on the server.
|
|
"""
|
|
opts = {k: v for k, v in context_options().items() if k != "process_revision_directives"}
|
|
diffs = await _run(
|
|
lambda c: compare_metadata(MigrationContext.configure(c, opts=opts), target_metadata())
|
|
)
|
|
if not diffs:
|
|
logger.info("schema matches the models")
|
|
return None
|
|
logger.info("%s schema difference(s) detected", len(diffs))
|
|
current_rev = await current()
|
|
if current_rev and not _revision_on_disk(current_rev):
|
|
logger.warning(
|
|
"database revision %s is not in migrations/versions/; "
|
|
"new revision will follow local head %s (stamp unchanged)",
|
|
current_rev,
|
|
head(),
|
|
)
|
|
return await _run(lambda c: _write_revision(c, message))
|
|
|
|
|
|
async def run_manual_sql() -> None:
|
|
"""Apply `migrations/manual/*.sql` once per database, in filename order.
|
|
|
|
This is what lets a developer just pull and boot: the one-shot data
|
|
migrations (enum labels, RBAC seed, backfills) apply themselves instead of
|
|
requiring a psql session. The files are written idempotent, but a tracking
|
|
table pins each to a single application per database; a file already run by
|
|
hand before this runner existed re-runs once (harmlessly) to get recorded.
|
|
|
|
Runs on the raw asyncpg connection: the files are multi-statement psql
|
|
batches, which the prepared-statement path cannot execute.
|
|
"""
|
|
files = sorted(p for p in (MIGRATIONS / "manual").glob("*.sql") if p.is_file())
|
|
if not files:
|
|
return
|
|
schema = get_settings().db_default_schema
|
|
table = f'"{schema}".{MANUAL_TABLE}' if schema else MANUAL_TABLE
|
|
async with get_engine().connect() as conn:
|
|
raw = await conn.get_raw_connection()
|
|
driver = raw.driver_connection
|
|
await driver.execute(
|
|
f"CREATE TABLE IF NOT EXISTS {table} ("
|
|
" filename text PRIMARY KEY,"
|
|
" applied_at timestamptz NOT NULL DEFAULT now())"
|
|
)
|
|
applied = {r["filename"] for r in await driver.fetch(f"SELECT filename FROM {table}")}
|
|
for path in files:
|
|
if path.name in applied:
|
|
continue
|
|
# utf-8-sig: Windows editors save SQL with a BOM, which would
|
|
# otherwise reach Postgres glued onto the first statement.
|
|
await driver.execute(path.read_text(encoding="utf-8-sig"))
|
|
await driver.execute(f"INSERT INTO {table} (filename) VALUES ($1)", path.name)
|
|
logger.info("applied manual migration %s", path.name)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _lock() -> AsyncIterator[None]:
|
|
"""Advisory lock, so only one worker migrates when several boot at once."""
|
|
async with get_engine().connect() as conn:
|
|
await conn.execution_options(isolation_level="AUTOCOMMIT")
|
|
await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": LOCK_ID})
|
|
try:
|
|
yield
|
|
finally:
|
|
await conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": LOCK_ID})
|
|
|
|
|
|
async def migrate(*, autogen: bool | None = None, message: str = "auto") -> None:
|
|
"""Bring the DB in line with models, then run manual SQL, under the lock.
|
|
|
|
Empty Compose volumes bootstrap from models (create_all + stamp). Every boot
|
|
upgrades any on-disk revisions (skipped when versions are absent from the
|
|
image), then — when DB_AUTOGENERATE is on — applies ORM drift in-memory so
|
|
no `versions/*.py` files are written on the server.
|
|
"""
|
|
should_autogen = get_settings().db_autogenerate if autogen is None else autogen
|
|
async with _lock():
|
|
if await _schema_is_empty():
|
|
await bootstrap_empty()
|
|
else:
|
|
await upgrade()
|
|
if should_autogen:
|
|
try:
|
|
await apply_model_drift()
|
|
except Exception as exc:
|
|
# Drift can still trip on unrelated tables; file revisions already
|
|
# ran above. Log and continue so the API can finish booting.
|
|
logger.exception("ORM drift apply failed; continuing boot: %s", exc)
|
|
await run_manual_sql()
|
|
logger.info("database at revision %s", await current())
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
|
parser = argparse.ArgumentParser(prog="alembic_setup.py", description=__doc__)
|
|
parser.add_argument(
|
|
"command",
|
|
nargs="?",
|
|
default="migrate",
|
|
choices=[
|
|
"migrate",
|
|
"revision",
|
|
"makemigrations",
|
|
"upgrade",
|
|
"downgrade",
|
|
"current",
|
|
"head",
|
|
"stamp",
|
|
],
|
|
)
|
|
parser.add_argument("-m", "--message", default="auto", help="revision message")
|
|
parser.add_argument("-r", "--revision", help="target revision")
|
|
args = parser.parse_args(argv)
|
|
logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")
|
|
|
|
async def run() -> None:
|
|
from db_setup import init_db # local import: db_setup imports this module lazily
|
|
|
|
try:
|
|
if args.command == "migrate":
|
|
await init_db()
|
|
elif args.command in ("revision", "makemigrations"):
|
|
try:
|
|
print(await autogenerate(args.message) or "no changes")
|
|
except RuntimeError as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
raise SystemExit(1) from None
|
|
elif args.command == "upgrade":
|
|
await upgrade(args.revision or "head")
|
|
elif args.command == "downgrade":
|
|
await downgrade(args.revision or "-1")
|
|
elif args.command == "current":
|
|
print(await current())
|
|
elif args.command == "head":
|
|
print(head())
|
|
elif args.command == "stamp":
|
|
await stamp(args.revision or "head")
|
|
finally:
|
|
await close_db()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|