86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
"""Serialize optional schema reset and Alembic upgrade under one DB lock."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
from alembic.config import Config
|
|
from sqlalchemy import create_engine, inspect, text
|
|
|
|
from alembic import command
|
|
from src.core.env_settings import database_url
|
|
|
|
_MIGRATION_LOCK_ID = 481920260824
|
|
_BASELINE_REVISION = "0001_baseline"
|
|
|
|
|
|
def _enabled(value: str | None) -> bool:
|
|
return str(value or "").strip().lower() in {"1", "true", "yes", "y"}
|
|
|
|
|
|
def _maybe_reset(connection) -> None:
|
|
reset_enabled = _enabled(os.getenv("RESET_DATABASE_SCHEMA"))
|
|
sys.stderr.write(
|
|
f"[database] RESET_DATABASE_SCHEMA={'enabled' if reset_enabled else 'disabled'}\n"
|
|
)
|
|
if not reset_enabled:
|
|
return
|
|
|
|
if connection.dialect.name != "postgresql":
|
|
raise RuntimeError("RESET_DATABASE_SCHEMA is supported only for PostgreSQL")
|
|
|
|
tables = set(inspect(connection).get_table_names())
|
|
revision = None
|
|
if "alembic_version" in tables:
|
|
revision = connection.execute(
|
|
text("SELECT version_num FROM alembic_version LIMIT 1")
|
|
).scalar()
|
|
|
|
if revision == _BASELINE_REVISION:
|
|
sys.stderr.write("[database] Baseline already applied; one-shot reset skipped\n")
|
|
connection.commit()
|
|
return
|
|
if tables and revision is None:
|
|
raise RuntimeError(
|
|
"Database schema is non-empty but has no Alembic revision; "
|
|
"refusing destructive reset"
|
|
)
|
|
|
|
connection.execute(text("DROP SCHEMA IF EXISTS public CASCADE"))
|
|
connection.execute(text("CREATE SCHEMA public"))
|
|
connection.commit()
|
|
sys.stderr.write(
|
|
f"[database] Legacy Alembic revision {revision or '<empty>'} reset successfully\n"
|
|
)
|
|
|
|
|
|
def prepare_database() -> None:
|
|
engine = create_engine(database_url(), pool_pre_ping=True)
|
|
try:
|
|
with engine.connect() as connection:
|
|
locked = connection.dialect.name == "postgresql"
|
|
if locked:
|
|
connection.execute(text("SELECT pg_advisory_lock(:lock_id)"), {"lock_id": _MIGRATION_LOCK_ID})
|
|
connection.commit()
|
|
try:
|
|
_maybe_reset(connection)
|
|
connection.commit()
|
|
config = Config(str(Path(__file__).resolve().parents[2] / "alembic.ini"))
|
|
config.attributes["connection"] = connection
|
|
command.upgrade(config, "head")
|
|
finally:
|
|
if locked:
|
|
connection.execute(
|
|
text("SELECT pg_advisory_unlock(:lock_id)"),
|
|
{"lock_id": _MIGRATION_LOCK_ID},
|
|
)
|
|
connection.commit()
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
prepare_database()
|