HR-ATS-Portal/backend/alembic_setup.py

"""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|upgrade|downgrade|current|head] [-m MSG] [-r REV]

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
from alembic.config import Config
from alembic.runtime.migration import MigrationContext
from alembic.script import ScriptDirectory
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_ != "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,
        "compare_server_default": True,
        "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."""
    async with get_engine().connect() as conn:
        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())


async def upgrade(revision: str = "head") -> None:
    await _run(lambda c: command.upgrade(config(c), revision))
    logger.info("upgraded to %s", revision)


async def downgrade(revision: str = "-1") -> None:
    await _run(lambda c: command.downgrade(config(c), revision))
    logger.info("downgraded to %s", revision)


async def autogenerate(message: str = "auto") -> str | None:
    """Write a revision if the models have drifted; return its id, or None."""
    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))
    before = head()
    await _run(lambda c: command.revision(config(c), message=message, autogenerate=True))
    after = head()
    return after if after != before else None


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:
    """Apply pending revisions, fresh model drift, then manual SQL, under the lock."""
    should_autogen = get_settings().db_autogenerate if autogen is None else autogen
    async with _lock():
        await upgrade()
        if should_autogen and await autogenerate(message):
            await upgrade()
        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", "upgrade", "downgrade", "current", "head"],
    )
    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 == "revision":
                print(await autogenerate(args.message) or "no changes")
            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())
        finally:
            await close_db()

    asyncio.run(run())


if __name__ == "__main__":
    main()