"""Alembic helpers for PortfolioHub database migrations.""" from __future__ import annotations from importlib import resources import psycopg from alembic import command from alembic.config import Config from alembic.script import ScriptDirectory from app.core.config import Settings def migration_root(): """Return the packaged Alembic resource tree for source and wheel installs.""" return resources.files("app.db.alembic") def migration_inventory() -> tuple[str, ...]: """Return the deterministic, non-empty migration revision inventory.""" versions = migration_root().joinpath("versions") inventory = tuple(sorted(item.name for item in versions.iterdir() if item.name.endswith(".py") and item.name != "__init__.py")) if not inventory: raise RuntimeError("packaged Alembic migration inventory is empty") return inventory def get_alembic_config(settings: Settings) -> Config: root = migration_root() config = Config(str(root.joinpath("alembic.ini"))) config.set_main_option("script_location", str(root)) config.set_main_option("sqlalchemy.url", settings.active_database_url.replace("%", "%%")) # migrations/env.py must use this explicit settings object instead of the # process-global cached Settings. This keeps programmatic callers (notably # isolated integration tests and runtime bootstrap) on their requested DB. config.attributes["portfoliohub_settings"] = settings return config def upgrade_to_head(settings: Settings) -> None: command.upgrade(get_alembic_config(settings), "head") def bootstrap_or_upgrade_to_head(settings: Settings) -> str: if _has_alembic_version_table(settings): upgrade_to_head(settings) return "head" if _has_legacy_portfolio_schema(settings): # Runtime ensure_schema() can create the BAT-025 columns before Alembic is # introduced. Stamp only the revision immediately before BAT-025, then run # the real upgrade so 0028 backfills a nullable pre-created `source` and # converges constraints/defaults without replaying the older DDL chain. config = get_alembic_config(settings) command.stamp(config, "20260628_0027") command.upgrade(config, "head") return "head" upgrade_to_head(settings) return "head" def current_revision(settings: Settings) -> str | None: """Read the database revision without applying or stamping migrations.""" with _connect(settings) as conn, conn.cursor() as cur: cur.execute("SELECT version_num FROM alembic_version") row = cur.fetchone() return str(row[0]) if row else None def head_revision(settings: Settings) -> str: """Return the single packaged Alembic head without touching the database.""" heads = ScriptDirectory.from_config(get_alembic_config(settings)).get_heads() if len(heads) != 1: raise RuntimeError(f"expected one Alembic head, found {len(heads)}") return heads[0] def _connect(settings: Settings) -> psycopg.Connection: return psycopg.connect( host=settings.active_db_host, port=settings.active_db_port, dbname=settings.active_db_name, user=settings.active_db_user, password=settings.active_db_password, sslmode=settings.active_db_sslmode, connect_timeout=settings.active_db_connect_timeout_sec, application_name=settings.active_db_application_name, autocommit=True, ) def _has_alembic_version_table(settings: Settings) -> bool: with _connect(settings) as conn, conn.cursor() as cur: cur.execute("SELECT to_regclass('public.alembic_version')") return cur.fetchone()[0] is not None def _has_legacy_portfolio_schema(settings: Settings) -> bool: required_tables = {"strategy_sets", "portfolios", "portfolio_ledger"} with _connect(settings) as conn: with conn.cursor() as cur: cur.execute( "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND tablename = ANY(%s)", (list(required_tables),), ) found = {row[0] for row in cur.fetchall()} return found == required_tables