calminer/scripts/backfill_currency.py

"""
Backfill script to populate currency_id for capex and opex rows using existing currency_code.

Usage:
  python scripts/backfill_currency.py --dry-run
  python scripts/backfill_currency.py --create-missing

This script is intentionally cautious: it defaults to dry-run mode and will refuse to run
if database connection settings are missing. It supports creating missing currency rows when `--create-missing`
is provided. Always run against a development/staging database first.
"""

from __future__ import annotations
import argparse
import importlib
import sys
from pathlib import Path

from sqlalchemy import text, create_engine


PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def load_database_url() -> str:
    try:
        db_module = importlib.import_module("config.database")
    except RuntimeError as exc:
        raise RuntimeError(
            "Database configuration missing: set DATABASE_URL or provide granular "
            "variables (DATABASE_DRIVER, DATABASE_HOST, DATABASE_PORT, DATABASE_USER, "
            "DATABASE_PASSWORD, DATABASE_NAME, optional DATABASE_SCHEMA)."
        ) from exc

    return getattr(db_module, "DATABASE_URL")


def backfill(
    db_url: str, dry_run: bool = True, create_missing: bool = False
) -> None:
    engine = create_engine(db_url)
    with engine.begin() as conn:
        # Ensure currency table exists
        res = (
            conn.execute(
                text(
                    "SELECT name FROM sqlite_master WHERE type='table' AND name='currency';"
                )
            )
            if db_url.startswith("sqlite:")
            else conn.execute(text("SELECT to_regclass('public.currency');"))
        )
        # Note: we don't strictly depend on the above - we assume migration was already applied

        # Helper: find or create currency by code
        def find_currency_id(code: str):
            r = conn.execute(
                text("SELECT id FROM currency WHERE code = :code"),
                {"code": code},
            ).fetchone()
            if r:
                return r[0]
            if create_missing:
                # insert and return id
                conn.execute(
                    text(
                        "INSERT INTO currency (code, name, symbol, is_active) VALUES (:c, :n, NULL, TRUE)"
                    ),
                    {"c": code, "n": code},
                )
                r2 = conn.execute(
                    text("SELECT id FROM currency WHERE code = :code"),
                    {"code": code},
                ).fetchone()
                if not r2:
                    raise RuntimeError(
                        f"Unable to determine currency ID for '{code}' after insert"
                    )
                return r2[0]
            return None

        # Process tables capex and opex
        for table in ("capex", "opex"):
            # Check if currency_id column exists
            try:
                cols = (
                    conn.execute(
                        text(
                            f"SELECT 1 FROM information_schema.columns WHERE table_name = '{table}' AND column_name = 'currency_id'"
                        )
                    )
                    if not db_url.startswith("sqlite:")
                    else [(1,)]
                )
            except Exception:
                cols = [(1,)]

            if not cols:
                print(f"Skipping {table}: no currency_id column found")
                continue

            # Find rows where currency_id IS NULL but currency_code exists
            rows = conn.execute(
                text(
                    f"SELECT id, currency_code FROM {table} WHERE currency_id IS NULL OR currency_id = ''"
                )
            )
            changed = 0
            for r in rows:
                rid = r[0]
                code = (r[1] or "USD").strip().upper()
                cid = find_currency_id(code)
                if cid is None:
                    print(
                        f"Row {table}:{rid} has unknown currency code '{code}' and create_missing=False; skipping"
                    )
                    continue
                if dry_run:
                    print(
                        f"[DRY RUN] Would set {table}.currency_id = {cid} for row id={rid} (code={code})"
                    )
                else:
                    conn.execute(
                        text(
                            f"UPDATE {table} SET currency_id = :cid WHERE id = :rid"
                        ),
                        {"cid": cid, "rid": rid},
                    )
                    changed += 1

            print(f"{table}: processed, changed={changed} (dry_run={dry_run})")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill currency_id from currency_code for capex/opex tables"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        default=True,
        help="Show actions without writing",
    )
    parser.add_argument(
        "--create-missing",
        action="store_true",
        help="Create missing currency rows in the currency table",
    )
    args = parser.parse_args()

    db = load_database_url()
    backfill(db, dry_run=args.dry_run, create_missing=args.create_missing)


if __name__ == "__main__":
    main()