procurement-ingest/scripts/migrate_tables.py

"""
One-off cross-account DynamoDB backfill for the mgmt -> seahaven-prod
migration (Phase 4 of the migration plan). Scans each table in the source
(mgmt) account and writes into the same-named, already-CDK-created table in
the destination (prod) account.

Dry-run is the DEFAULT in every mode; nothing is written or deleted unless
--execute is passed.

Per-table modes (encode the plan's dual-delivery-safe semantics):

  WorkOrders, purchase-orders   overwrite Put. Safe because during dual
                                delivery the mgmt item is a strict SUPERSET of
                                the prod item (both pipelines apply SET-only
                                merges to the same inbound mail).
  WorkOrderComments             cutoff-filtered Put: only rows with
                                attribute_not_exists(ingested_at) OR
                                ingested_at < --cutoff (T_activate). The same
                                email yields DIFFERENT comment_ids per account
                                (S3-object-key hash suffix), so copying
                                dual-window rows would create duplicate
                                history. --cutoff is REQUIRED for this table.
  verified-sites,               truncate-and-load: delete ALL destination rows
  pending-site-review           first, then copy the source scan verbatim.
                                Required because the purchase-orders backfill
                                replays the prod site-extractor stream, which
                                creates junk `pending` rows and non-idempotent
                                poCount ADDs; mgmt values are ground truth.
                                Take an on-demand backup of the dest table and
                                wait for it to be AVAILABLE before running
                                (the script refuses without --backup-arn).

Usage:
    # copy (dry-run first, then --execute)
    python scripts/migrate_tables.py copy --table WorkOrders
    python scripts/migrate_tables.py copy --table WorkOrderComments \
        --cutoff 2026-07-24T02:00:00Z --execute
    python scripts/migrate_tables.py copy --table verified-sites \
        --backup-arn arn:aws:dynamodb:...:backup/... --execute
    # verify (count parity + spot checks; always read-only)
    python scripts/migrate_tables.py verify --table purchase-orders
    python scripts/migrate_tables.py verify --all

Profiles: --source-profile (default: "default", the mgmt SSO session) and
--dest-profile (default: "seahaven-prod"). The script hard-verifies both
resolved account ids before doing anything.
"""

import argparse
import json
import sys

import boto3

SOURCE_ACCOUNT = "328440206208"
DEST_ACCOUNT = "011934824531"
REGION = "us-east-1"

TABLES = {
    "purchase-orders": {"mode": "overwrite", "keys": ["po_number"]},
    "WorkOrders": {"mode": "overwrite", "keys": ["work_order_id"]},
    "WorkOrderComments": {
        "mode": "cutoff",
        "keys": ["work_order_id", "comment_id"],
        "cutoff_attr": "ingested_at",
    },
    "verified-sites": {"mode": "truncate_load", "keys": ["siteCode"]},
    "pending-site-review": {"mode": "truncate_load", "keys": ["po_number"]},
}


def session_table(profile, expected_account, table_name):
    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.resource("dynamodb").Table(table_name)


def scan_items(table, filter_kwargs=None):
    kwargs = dict(filter_kwargs or {})
    while True:
        page = table.scan(**kwargs)
        yield from page.get("Items", [])
        lek = page.get("LastEvaluatedKey")
        if not lek:
            return
        kwargs["ExclusiveStartKey"] = lek


def scan_count(table, filter_kwargs=None):
    kwargs = dict(filter_kwargs or {}, Select="COUNT")
    total = 0
    while True:
        page = table.scan(**kwargs)
        total += page["Count"]
        lek = page.get("LastEvaluatedKey")
        if not lek:
            return total
        kwargs["ExclusiveStartKey"] = lek


def comments_cutoff_filter(cutoff):
    return {
        "FilterExpression": ("attribute_not_exists(#ia) OR #ia < :cutoff"),
        "ExpressionAttributeNames": {"#ia": "ingested_at"},
        "ExpressionAttributeValues": {":cutoff": cutoff},
    }


def normalize_cutoff(value):
    # ingested_at is datetime.now(timezone.utc).isoformat() -> "+00:00"
    # offset form. Normalize a trailing Z so string comparison is apples
    # to apples.
    return value.replace("Z", "+00:00")


def _truncate_destination(name, cfg, dst, args):
    """Backup-gated destination truncate for the truncate-and-load tables."""
    if args.execute:
        if not args.backup_arn:
            sys.exit(
                f"ERROR: {name} is truncate-and-load; pass --backup-arn of an "
                "AVAILABLE on-demand backup of the DESTINATION table taken "
                "just now (aws dynamodb create-backup / describe-backup)."
            )
        desc = dst.meta.client.describe_backup(BackupArn=args.backup_arn)
        details = desc["BackupDescription"]["BackupDetails"]
        src_table = desc["BackupDescription"]["SourceTableDetails"]
        if details["BackupStatus"] != "AVAILABLE" or src_table["TableName"] != name:
            sys.exit(
                f"ERROR: backup {args.backup_arn} is "
                f"{details['BackupStatus']} for table "
                f"{src_table['TableName']!r}; need an AVAILABLE backup "
                f"of {name!r}. Aborting before truncate."
            )
    dest_keys = [{k: item[k] for k in cfg["keys"]} for item in scan_items(dst)]
    print(f"[{name}] truncate: {len(dest_keys)} destination rows to delete")
    if args.execute:
        with dst.batch_writer() as batch:
            for key in dest_keys:
                batch.delete_item(Key=key)
        print(f"[{name}] truncate: done")
    else:
        print(f"[{name}] DRY-RUN: no deletes performed")


def copy_table(name, cfg, args):
    src = session_table(args.source_profile, SOURCE_ACCOUNT, name)
    dst = session_table(args.dest_profile, DEST_ACCOUNT, name)

    filter_kwargs = None
    if cfg["mode"] == "cutoff":
        if not args.cutoff:
            sys.exit(f"ERROR: --cutoff (T_activate) is required for {name}.")
        filter_kwargs = comments_cutoff_filter(normalize_cutoff(args.cutoff))

    if cfg["mode"] == "truncate_load":
        _truncate_destination(name, cfg, dst, args)

    copied = 0
    if args.execute:
        with dst.batch_writer() as batch:
            for item in scan_items(src, filter_kwargs):
                batch.put_item(Item=item)
                copied += 1
    else:
        for _ in scan_items(src, filter_kwargs):
            copied += 1
    verb = "copied" if args.execute else "would copy (DRY-RUN)"
    print(f"[{name}] {verb} {copied} items {SOURCE_ACCOUNT} -> {DEST_ACCOUNT}")


def verify_table(name, cfg, args):
    src = session_table(args.source_profile, SOURCE_ACCOUNT, name)
    dst = session_table(args.dest_profile, DEST_ACCOUNT, name)

    src_filter = None
    if cfg["mode"] == "cutoff" and args.cutoff:
        src_filter = comments_cutoff_filter(normalize_cutoff(args.cutoff))

    src_count = scan_count(src, src_filter)
    dst_count = scan_count(dst)
    # Destination may legitimately EXCEED source-side-of-cutoff (prod-native
    # dual-delivery rows) for cutoff tables, and match-or-exceed for overwrite
    # tables written after the copy; exact parity is expected only for
    # truncate_load immediately after the load.
    relation = {
        "overwrite": dst_count >= src_count,
        "cutoff": dst_count >= src_count,
        "truncate_load": dst_count == src_count,
    }[cfg["mode"]]
    status = "OK" if relation else "MISMATCH"
    print(f"[{name}] source={src_count} dest={dst_count} -> {status}")

    ok = relation
    if args.sample and cfg["mode"] != "truncate_load":
        sampled = mismatched = 0
        for item in scan_items(src, src_filter):
            if sampled >= args.sample:
                break
            sampled += 1
            key = {k: item[k] for k in cfg["keys"]}
            got = dst.get_item(Key=key).get("Item")
            if got is None:
                mismatched += 1
                print(f"[{name}] MISSING in dest: {json.dumps(key, default=str)}")
        print(f"[{name}] spot-check: {sampled} sampled, {mismatched} missing")
        ok = ok and mismatched == 0

    cancelled_check_limit = 25
    if name == "purchase-orders":
        # Sticky-Cancelled spot check: every source-Cancelled PO must still be
        # Cancelled in dest (the overwrite Put must not have raced a revival).
        cancelled_filter = {
            "FilterExpression": "po_status = :c",
            "ExpressionAttributeValues": {":c": "Cancelled"},
        }
        checked = wrong = 0
        for item in scan_items(src, cancelled_filter):
            if checked >= cancelled_check_limit:
                break
            checked += 1
            got = dst.get_item(Key={"po_number": item["po_number"]}).get("Item")
            if not got or got.get("po_status") != "Cancelled":
                wrong += 1
                print(f"[{name}] CANCELLED-DRIFT: {item['po_number']}")
        print(f"[{name}] cancelled spot-check: {checked} checked, {wrong} drifted")
        ok = ok and wrong == 0

    return ok


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    sub = ap.add_subparsers(dest="cmd", required=True)
    for cmd in ("copy", "verify"):
        p = sub.add_parser(cmd)
        p.add_argument("--table", choices=sorted(TABLES))
        p.add_argument("--all", action="store_true", help="all five tables")
        p.add_argument("--source-profile", default="default")
        p.add_argument("--dest-profile", default="seahaven-prod")
        p.add_argument(
            "--cutoff",
            help="T_activate ISO-8601 UTC; required to copy WorkOrderComments",
        )
        if cmd == "copy":
            p.add_argument("--execute", action="store_true")
            p.add_argument(
                "--backup-arn",
                help="required with --execute for truncate-and-load tables",
            )
        else:
            p.add_argument(
                "--sample",
                type=int,
                default=10,
                help="per-table spot-check sample size (0 disables)",
            )

    args = ap.parse_args()
    if bool(args.table) == bool(args.all):
        sys.exit("ERROR: pass exactly one of --table or --all.")
    names = sorted(TABLES) if args.all else [args.table]

    if args.cmd == "copy":
        # Plan-mandated order: WorkOrders, WorkOrderComments, purchase-orders
        # (stream replay), drain, then the two truncate-and-load tables. With
        # --all we enforce it; single --table trusts the operator.
        order = [
            "WorkOrders",
            "WorkOrderComments",
            "purchase-orders",
            "verified-sites",
            "pending-site-review",
        ]
        for name in sorted(names, key=order.index):
            if TABLES[name]["mode"] == "truncate_load" and args.all:
                sys.exit(
                    "ERROR: run verified-sites / pending-site-review "
                    "individually AFTER the purchase-orders stream drain "
                    "(ESM IteratorAge ~ 0), not as part of --all."
                )
            copy_table(name, TABLES[name], args)
    else:
        results = [verify_table(n, TABLES[n], args) for n in names]
        if not all(results):
            sys.exit(1)


if __name__ == "__main__":
    main()