afterhours-shift-manager/scripts/cutover/copy_dynamodb.py

#!/usr/bin/env python3
"""Copy afterhours-shifts from mgmt to prod. Dry-run unless --execute."""

from __future__ import annotations

import argparse
import sys

import boto3


TABLE = "afterhours-shifts"
SRC_ACCOUNT = "328440206208"
DST_ACCOUNT = "011934824531"


def _client(profile: str, region: str):
    session = boto3.Session(profile_name=profile, region_name=region)
    return session.client("dynamodb")


def _scan_all(client):
    items = []
    kwargs = {"TableName": TABLE}
    while True:
        resp = client.scan(**kwargs)
        items.extend(resp.get("Items", []))
        start = resp.get("LastEvaluatedKey")
        if not start:
            return items
        kwargs["ExclusiveStartKey"] = start


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src-profile", required=True)
    parser.add_argument("--dst-profile", required=True)
    parser.add_argument("--region", default="us-east-1")
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()

    src = _client(args.src_profile, args.region)
    dst = _client(args.dst_profile, args.region)
    src_id = boto3.Session(profile_name=args.src_profile).client("sts").get_caller_identity()["Account"]
    dst_id = boto3.Session(profile_name=args.dst_profile).client("sts").get_caller_identity()["Account"]
    if src_id != SRC_ACCOUNT:
        print(f"src account {src_id} is not mgmt {SRC_ACCOUNT}", file=sys.stderr)
        return 2
    if dst_id != DST_ACCOUNT:
        print(f"dst account {dst_id} is not prod {DST_ACCOUNT}", file=sys.stderr)
        return 2

    items = _scan_all(src)
    dst_count = dst.describe_table(TableName=TABLE)["Table"]["ItemCount"]
    print(f"src items={len(items)} dst describe ItemCount={dst_count}")
    if not args.execute:
        print("dry-run; pass --execute to BatchWriteItem")
        return 0

    written = 0
    batch = []
    for item in items:
        batch.append({"PutRequest": {"Item": item}})
        if len(batch) == 25:
            dst.batch_write_item(RequestItems={TABLE: batch})
            written += len(batch)
            batch = []
    if batch:
        dst.batch_write_item(RequestItems={TABLE: batch})
        written += len(batch)
    after = _scan_all(dst)
    print(f"wrote={written} dst_scan={len(after)}")
    if len(after) != len(items):
        print("item counts differ after copy", file=sys.stderr)
        return 1
    return 0


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