DataOps-platform/scripts/create_metadata_review_tables.py

#!/usr/bin/env python3
"""
Create Metadata Review tables in PostgreSQL (production by default).

Tables:
  - public.metadata_review_records
  - public.metadata_version_history

This script executes the SQL files under /database:
  - database/create_metadata_review_records_table.sql
  - database/create_metadata_version_history_table.sql

Usage:
  python scripts/create_metadata_review_tables.py
  python scripts/create_metadata_review_tables.py --db-uri postgresql://...
  python scripts/create_metadata_review_tables.py --env production
  python scripts/create_metadata_review_tables.py --dry-run
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path
from typing import Optional


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("CreateMetadataReviewTables")

WORKSPACE_ROOT = Path(__file__).resolve().parent.parent
SQL_REVIEW = WORKSPACE_ROOT / "database" / "create_metadata_review_records_table.sql"
SQL_HISTORY = (
    WORKSPACE_ROOT / "database" / "create_metadata_version_history_table.sql"
)


def load_db_uri(env: str, override_uri: Optional[str]) -> str:
    if override_uri:
        return override_uri

    sys.path.insert(0, str(WORKSPACE_ROOT))
    from app.config.config import config  # type: ignore

    if env not in config:
        raise ValueError(f"Unknown env: {env}. Must be one of: {list(config.keys())}")

    return config[env].SQLALCHEMY_DATABASE_URI


def execute_sql_file(db_uri: str, sql_path: Path, dry_run: bool) -> None:
    if not sql_path.exists():
        raise FileNotFoundError(f"SQL file not found: {sql_path}")

    sql = sql_path.read_text(encoding="utf-8")
    logger.info(f"Applying SQL: {sql_path.relative_to(WORKSPACE_ROOT)}")

    if dry_run:
        logger.info("dry-run enabled; skipping execution.")
        return

    import psycopg2  # type: ignore

    conn = psycopg2.connect(db_uri)
    try:
        with conn:
            with conn.cursor() as cur:
                cur.execute(sql)
        logger.info("OK")
    finally:
        conn.close()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Create metadata review tables in PostgreSQL",
    )
    parser.add_argument(
        "--env",
        default="production",
        help="Config env name in app.config.config (default: production)",
    )
    parser.add_argument(
        "--db-uri",
        default=None,
        help="Override database URI (postgresql://...)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print actions without executing SQL",
    )
    args = parser.parse_args()

    db_uri = load_db_uri(args.env, args.db_uri)
    safe_target = db_uri.split("@")[-1] if "@" in db_uri else db_uri
    logger.info(f"Target database: {safe_target}")

    execute_sql_file(db_uri, SQL_REVIEW, dry_run=bool(args.dry_run))
    execute_sql_file(db_uri, SQL_HISTORY, dry_run=bool(args.dry_run))
    logger.info("All done.")


if __name__ == "__main__":
    main()