Source code for trader.persistence.migrations

"""Alembic integration: stamping fresh databases, and upgrading existing ones.

**Why `create_all` still exists alongside migrations.** `init_db` builds a new
database with `create_all` rather than by replaying every migration, because
the test suite creates over a hundred throwaway databases and replaying the
history for each one costs far more than it proves. The risk that buys — the
two definitions drifting apart — is closed by
`tests/persistence/test_migrations.py`, which fails if the models and the
migrations disagree. That is this project's equivalent of Django's
`makemigrations --check`.

A database built by `create_all` is stamped at head immediately, so Alembic
knows its tables already exist and the next `upgrade` applies only genuinely
new revisions.
"""

from pathlib import Path

from alembic import command
from alembic.config import Config

__all__ = ["alembic_config", "current_revision", "stamp_head", "upgrade_to_head"]

#: Repo root: this file is `src/trader/persistence/migrations.py`.
_PROJECT_ROOT = Path(__file__).resolve().parents[3]
_ALEMBIC_INI = _PROJECT_ROOT / "alembic.ini"




def alembic_config(database_url: str) -> Config:
    """An Alembic config pointed at `database_url`.

    The URL is set here rather than read from `alembic.ini`, for the same
    reason `env.py` resolves it from `Settings`: one place decides where the
    database lives.
    """
    config = Config(str(_ALEMBIC_INI))
    config.set_main_option("script_location", str(_PROJECT_ROOT / "migrations"))
    config.set_main_option("sqlalchemy.url", database_url)
    return config





def stamp_head(database_url: str) -> None:
    """Record `database_url` as being at the latest revision, creating nothing.

    Used after `create_all` on a fresh database, and it is also how an existing
    database adopts Alembic. Running `upgrade` on either would fail on the
    first `CREATE TABLE`, because the tables are already there.
    """
    command.stamp(alembic_config(database_url), "head")





def upgrade_to_head(database_url: str) -> None:
    """Apply every migration `database_url` has not yet seen."""
    command.upgrade(alembic_config(database_url), "head")



def head_revision(database_url: str) -> str | None:
    """The revision the migration scripts define as head.

    Takes a URL only because `alembic_config` does; the answer depends on the
    scripts on disk, not on the database. Used to refuse to start an
    unattended daemon against a schema the code no longer matches — a
    half-migrated database that a human would notice immediately can otherwise
    fail quietly every fifteen minutes for a weekend.
    """
    from alembic.script import ScriptDirectory

    return ScriptDirectory.from_config(alembic_config(database_url)).get_current_head()




def current_revision(database_url: str) -> str | None:
    """The revision `database_url` is stamped at, or `None` if unstamped."""
    from alembic.runtime.migration import MigrationContext
    from sqlalchemy import create_engine

    engine = create_engine(database_url)
    try:
        with engine.connect() as connection:
            return MigrationContext.configure(connection).get_current_revision()
    finally:
        engine.dispose()