109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
"""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 |