procurement-ingest/scripts/migrate_wo_tables.py

#!/usr/bin/env python3
"""PLAT-11 same-account WO table copy: PascalCase -> kebab-case.

Copies ``WorkOrders`` -> ``work-orders`` and ``WorkOrderComments`` ->
``work-order-comments`` inside seahaven-prod. Dry-run is the DEFAULT; nothing
is written unless ``--execute`` is passed.

Usage:
    python scripts/migrate_wo_tables.py copy              # dry-run both
    python scripts/migrate_wo_tables.py copy --execute
    python scripts/migrate_wo_tables.py verify
    python scripts/migrate_wo_tables.py copy --table work-orders --execute
"""

from __future__ import annotations

import argparse
import sys
import time

import boto3
from boto3.dynamodb.types import TypeDeserializer

EXPECTED_ACCOUNT = "011934824531"
REGION = "us-east-1"

PAIRS = {
    "work-orders": {
        "source": "WorkOrders",
        "dest": "work-orders",
        "keys": ["work_order_id"],
    },
    "work-order-comments": {
        "source": "WorkOrderComments",
        "dest": "work-order-comments",
        "keys": ["work_order_id", "comment_id"],
    },
}

_deserializer = TypeDeserializer()


def _client(profile: str):
    session = boto3.Session(profile_name=profile, region_name=REGION)
    acct = session.client("sts").get_caller_identity()["Account"]
    if acct != EXPECTED_ACCOUNT:
        sys.exit(
            f"ERROR: profile {profile!r} resolves to account {acct}, "
            f"expected {EXPECTED_ACCOUNT}. Aborting."
        )
    return session.client("dynamodb")


def _describe_count(client, table: str) -> int:
    return int(client.describe_table(TableName=table)["Table"]["ItemCount"])


def _scan_all(client, table: str):
    kwargs = {"TableName": table}
    while True:
        resp = client.scan(**kwargs)
        for item in resp.get("Items", []):
            yield item
        if "LastEvaluatedKey" not in resp:
            break
        kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"]


def _batch_write(client, table: str, items: list[dict], execute: bool) -> int:
    if not execute:
        return len(items)
    written = 0
    for i in range(0, len(items), 25):
        chunk = items[i : i + 25]
        request_items = {
            table: [{"PutRequest": {"Item": item}} for item in chunk],
        }
        unprocessed = request_items
        while unprocessed:
            resp = client.batch_write_item(RequestItems=unprocessed)
            unprocessed = resp.get("UnprocessedItems") or {}
            if unprocessed:
                time.sleep(0.2)
        written += len(chunk)
    return written


def copy_table(client, pair_key: str, execute: bool) -> None:
    pair = PAIRS[pair_key]
    source, dest = pair["source"], pair["dest"]
    print(f"==> copy {source} -> {dest} ({'EXECUTE' if execute else 'dry-run'})")
    items = list(_scan_all(client, source))
    print(f"    scanned {len(items)} items from {source}")
    written = _batch_write(client, dest, items, execute=execute)
    print(f"    {'wrote' if execute else 'would write'} {written} items to {dest}")


def verify_table(client, pair_key: str) -> bool:
    pair = PAIRS[pair_key]
    source, dest = pair["source"], pair["dest"]
    # ItemCount is eventually consistent; prefer live scan counts for gate.
    src_items = list(_scan_all(client, source))
    dst_items = list(_scan_all(client, dest))
    src_n, dst_n = len(src_items), len(dst_items)
    print(f"==> verify {source} ({src_n}) vs {dest} ({dst_n})")
    if src_n != dst_n:
        print("    FAIL count mismatch")
        return False

    key_names = pair["keys"]

    def key_tuple(item):
        return tuple(_deserializer.deserialize(item[k]) for k in key_names)

    src_by_key = {key_tuple(i): i for i in src_items}
    dst_by_key = {key_tuple(i): i for i in dst_items}
    missing = [k for k in src_by_key if k not in dst_by_key]
    if missing:
        print(f"    FAIL missing keys on dest (showing up to 5): {missing[:5]}")
        return False

    # Spot-check up to 25 items for full attribute equality.
    checked = 0
    for key, src_item in list(src_by_key.items())[:25]:
        if src_item != dst_by_key[key]:
            print(f"    FAIL attribute mismatch for key {key}")
            return False
        checked += 1
    print(f"    OK counts equal; spot-checked {checked} items")
    return True


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["copy", "verify"])
    parser.add_argument(
        "--table",
        choices=list(PAIRS.keys()),
        help="Single pair key (default: both)",
    )
    parser.add_argument(
        "--execute",
        action="store_true",
        help="Actually write (copy only). Default is dry-run.",
    )
    parser.add_argument("--profile", default="seahaven-prod")
    args = parser.parse_args()

    client = _client(args.profile)
    keys = [args.table] if args.table else list(PAIRS.keys())

    if args.command == "copy":
        for key in keys:
            copy_table(client, key, execute=args.execute)
        if not args.execute:
            print("Dry-run complete. Re-run with --execute to write.")
        return 0

    ok = True
    for key in keys:
        ok = verify_table(client, key) and ok
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())