chore(release): add alembic preflight check

2026-06-21 13:01:49 +08:00
parent ccb6786e0b
commit e98dfc9ffb
3 changed files with 136 additions and 1 deletions
@@ -0,0 +1,84 @@
+"""Preflight Alembic state before a production release."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from dataclasses import dataclass
+
+from sqlalchemy import create_engine, text
+
+
+@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("alembic.ini")
+    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())