Xiaoxia AI
89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
"""Preflight Alembic state before a production release."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import create_engine, text
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AlembicPreflightResult:
|
|
has_alembic_version: bool
|
|
business_table_count: int
|
|
current_revision: str | None
|
|
head_revision: str
|
|
|
|
@property
|
|
def action(self) -> str:
|
|
if self.has_alembic_version:
|
|
if self.current_revision == self.head_revision:
|
|
return "upgrade_head_noop_or_verify"
|
|
return "upgrade_head"
|
|
if self.business_table_count > 0:
|
|
return "stamp_head_then_upgrade_head"
|
|
return "upgrade_head_empty_database"
|
|
|
|
|
|
def get_head_revision() -> str:
|
|
from alembic.config import Config
|
|
from alembic.script import ScriptDirectory
|
|
|
|
config = Config(str(REPO_ROOT / "alembic.ini"))
|
|
config.set_main_option("script_location", str(REPO_ROOT / "alembic"))
|
|
script = ScriptDirectory.from_config(config)
|
|
head = script.get_current_head()
|
|
if not head:
|
|
raise RuntimeError("Alembic head revision not found")
|
|
return head
|
|
|
|
|
|
def inspect_database(database_url: str, head_revision: str) -> AlembicPreflightResult:
|
|
engine = create_engine(database_url)
|
|
with engine.begin() as connection:
|
|
has_alembic_version = (
|
|
connection.execute(
|
|
text("SELECT to_regclass(:table_name)"), {"table_name": "public.alembic_version"}
|
|
).scalar()
|
|
is not None
|
|
)
|
|
business_table_count = connection.execute(
|
|
text("SELECT COUNT(*) FROM pg_tables " "WHERE schemaname = :schema_name AND tablename != :version_table"),
|
|
{"schema_name": "public", "version_table": "alembic_version"},
|
|
).scalar_one()
|
|
current_revision = None
|
|
if has_alembic_version:
|
|
current_revision = connection.execute(text("SELECT version_num FROM alembic_version LIMIT 1")).scalar()
|
|
|
|
return AlembicPreflightResult(
|
|
has_alembic_version=has_alembic_version,
|
|
business_table_count=int(business_table_count),
|
|
current_revision=current_revision,
|
|
head_revision=head_revision,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--database-url", default=os.getenv("DATABASE_URL"), help="Defaults to DATABASE_URL.")
|
|
args = parser.parse_args()
|
|
if not args.database_url:
|
|
raise SystemExit("DATABASE_URL is required")
|
|
|
|
result = inspect_database(args.database_url, get_head_revision())
|
|
print(f"has_alembic_version={str(result.has_alembic_version).lower()}")
|
|
print(f"business_table_count={result.business_table_count}")
|
|
print(f"current_revision={result.current_revision or '<none>'}")
|
|
print(f"head_revision={result.head_revision}")
|
|
print(f"recommended_action={result.action}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|