diff --git a/tools/ddb_migration/.gitignore b/tools/ddb_migration/.gitignore new file mode 100644 index 00000000..708b79a1 --- /dev/null +++ b/tools/ddb_migration/.gitignore @@ -0,0 +1,15 @@ +.DS_Store +__pycache__/ +*.py[cod] +.venv +venv +.pytest_cache +.coverage +.coverage.* +htmlcov/ +coverage.xml + +# Demo artifacts +demo/config.env +demo/.demo_state +*.zip diff --git a/tools/ddb_migration/AGENTS.md b/tools/ddb_migration/AGENTS.md new file mode 100644 index 00000000..8395e456 --- /dev/null +++ b/tools/ddb_migration/AGENTS.md @@ -0,0 +1,45 @@ +# Agents Guide — `tools/ddb_migration` + +Conventions for future contributors (human or AI) extending this toolkit. + +## Module boundaries + +- **`transform.py`** — single source of truth for any per-item transform. Both `lambda/stream_replay.py` and `scripts/backfill.py` import it. Keep it free of AWS SDK calls. +- **`lambda/stream_replay.py`** — runs in Lambda. No filesystem writes, no `print()` for structured data (use the `_log()` helper). Every record write must be conditional on `_migration_ts`. +- **`scripts/backfill.py`** — runs locally or in EC2/CodeBuild. Same conditional-write contract as the Lambda. Always sets `_migration_ts = 0`. +- **`scripts/convergence_check.py`** — gate (exits 0 / 1). Do not add interactive prompts. +- **`scripts/cleanup.py`** — idempotent post-cutover. Safe to re-run. +- **`scripts/verify_cutover.py`** — read-only. Never writes to source or target. + +## Conditional-write contract (load-bearing) + +Every write to the target table goes through: + +```python +table.put_item( + Item=item, + ConditionExpression='attribute_not_exists(#pk) OR #ts < :ts', + ExpressionAttributeNames={'#pk': partition_key, '#ts': '_migration_ts'}, + ExpressionAttributeValues={':ts': migration_ts}, +) +``` + +`migration_ts` rules: +- backfill writes: `0` +- stream replay `INSERT`/`MODIFY`: `event['dynamodb']['ApproximateCreationDateTime']` +- stream replay `REMOVE` (tombstone): `event['dynamodb']['ApproximateCreationDateTime']` + +Do not bypass this. Do not introduce a third writer with a different timestamp source. + +## Adding a new script + +1. Add `scripts/your_script.py` with a `main()` that returns an `int` exit code. +2. Add `tests/test_your_script.py` using moto fixtures from `conftest.py`. +3. Document in `README.md` under "Scripts." +4. If it provisions AWS resources, hook it into `deploy.sh` AND `teardown.sh`. + +## Testing + +- Unit tests use `moto`. No real AWS calls. +- DynamoDB Streams events are constructed as fixtures (moto's stream support is partial); we mock the Lambda handler's boto3 client when needed. +- `make test` must stay green before any commit. diff --git a/tools/ddb_migration/Makefile b/tools/ddb_migration/Makefile new file mode 100644 index 00000000..8ec56103 --- /dev/null +++ b/tools/ddb_migration/Makefile @@ -0,0 +1,44 @@ +.DEFAULT_GOAL := help +.PHONY: help install test coverage lint clean + +VENV := .venv +PIP := $(VENV)/bin/pip +PYTEST := $(VENV)/bin/pytest + +PYTHON := $(shell command -v python3.11 || command -v python3.10 || command -v python3) + +help: + @echo "Targets:" + @echo " make install Create venv and install dev dependencies" + @echo " make test Run unit tests with coverage summary" + @echo " make coverage Run tests with detailed missing-line coverage" + @echo " make lint Syntax check shell + python sources" + @echo " make clean Remove venv and caches" + +install: + $(PYTHON) -m venv $(VENV) + $(PIP) install --upgrade pip + $(PIP) install -r requirements-dev.txt + +test: + @$(PYTEST) -q --tb=short --cov=scripts --cov=lambda --cov-branch --cov-report= 2>&1 | tail -10 + @$(VENV)/bin/coverage json -o .coverage.json --quiet 2>/dev/null && \ + $(VENV)/bin/python -c "import json; t=json.load(open('.coverage.json'))['totals']; \ + print(); \ + print(f\" LINE COVERAGE: {t['percent_statements_covered']:5.1f}% ({t['covered_lines']}/{t['num_statements']} statements covered)\"); \ + print(f\" BRANCH COVERAGE: {t['percent_branches_covered']:5.1f}% ({t['covered_branches']}/{t['num_branches']} branches covered)\"); \ + print()" && rm -f .coverage.json + +coverage: + $(PYTEST) --cov=scripts --cov=lambda --cov-branch --cov-report=term-missing + +lint: + @for f in deploy.sh teardown.sh demo/run_demo.sh; do bash -n $$f && echo " OK $$f" || exit 1; done + @find . -name '*.py' -not -path './.venv/*' -exec $(VENV)/bin/python -m py_compile {} + + @echo " OK python syntax" + +clean: + rm -rf $(VENV) .pytest_cache + find . -type d -name __pycache__ -exec rm -rf {} + + find . -type f -name .coverage -delete + find . -type f -name .coverage.json -delete diff --git a/tools/ddb_migration/NOTICE b/tools/ddb_migration/NOTICE new file mode 100644 index 00000000..7e3a7764 --- /dev/null +++ b/tools/ddb_migration/NOTICE @@ -0,0 +1,2 @@ +DynamoDB Zero-Downtime Migration Toolkit +Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. diff --git a/tools/ddb_migration/README.md b/tools/ddb_migration/README.md new file mode 100644 index 00000000..d1999bb0 --- /dev/null +++ b/tools/ddb_migration/README.md @@ -0,0 +1,210 @@ +# DynamoDB Zero-Downtime Migration Toolkit + +General-purpose, production-grade tooling for migrating an Amazon DynamoDB +table to a new table without taking writes offline. Built on three native +DynamoDB features — Export to S3, Streams, and conditional writes — and a +single-attribute conflict-resolution scheme. + +Use cases: + +* Adopting Multi-Region Strong Consistency (MRSC) on Global Tables (requires + starting from an empty replica). +* Cross-account migrations (workload consolidation, account splits, compliance). +* Schema changes — partition→composite key, attribute renames, new GSI keys. +* Billing-mode transitions on large tables without surprise throttling. +* Single-table consolidation across small tables. + +## Architecture + +Three overlapping phases: + +1. **Capture & bulk copy** — enable Streams on the source, export it to S3, + start the stream-replay Lambda, run `backfill.py` to load the export into + the target. Backfill writes go in with `_migration_ts=0`. +2. **Catch-up** — the Lambda replays every live source-table change to the + target, with `_migration_ts = ApproximateCreationDateTime` (always > 0). + The conditional expression `attribute_not_exists(#pk) OR #ts < :ts` ensures + newer timestamps always win, so backfill writes never overwrite live writes + and stale events never overwrite newer ones. +3. **Convergence & cutover** — `convergence_check.py` blocks until iterator age + is near zero, the DLQ is empty, and Scan COUNT on both tables agrees. Then + you flip routing. + +REMOVE events are written as `_tombstone=True` items (not deletes) so the +in-flight backfill cannot resurrect them. Post-cutover, `cleanup.py` enables +DynamoDB TTL to expire them automatically. + +## Layout + +``` +tools/ddb_migration/ +├── deploy.sh / teardown.sh One-command provision / clean-up +├── transform.py Shared per-item transform (customize here) +├── lambda/stream_replay.py Streams → target, conditional writes +├── scripts/ +│ ├── backfill.py S3 export → target, parallel + throttled +│ ├── convergence_check.py Pre-cutover gate (exit 0 / 1) +│ ├── cleanup.py Post-cutover removal of migration metadata +│ └── verify_cutover.py Sample-based source ↔ target verifier +├── iam/policies.json Reference IAM policy templates +├── demo/ One-click demo (real AWS resources) +└── tests/ pytest unit tests with moto +``` + +## Prerequisites + +* Python 3.10+ and `pip install -r requirements.txt`. +* AWS CLI v2, configured for the source-table account. +* PITR enabled on the source table (for the export). `deploy.sh` does not + enable this for you on production tables — verify before running. +* Permissions to create IAM roles, Lambda functions, S3 buckets, SQS queues, + SNS topics, and CloudWatch alarms. + +## Quick start (same account) + +```sh +cd tools/ddb_migration +make install +source .venv/bin/activate + +export SOURCE_TABLE=my-prod-table +export TARGET_TABLE=my-prod-table-v2 +export PARTITION_KEY=customer_id +export SORT_KEY=order_id # optional +export REGION=us-east-1 + +# 1. Provision Lambda, IAM role, DLQ, SNS topic, alarms, target table. +./deploy.sh + +# 2. Trigger an S3 export (the deploy.sh output prints the exact command). + +# 3. Run the backfill once the export completes. +EXPORT_BUCKET=ddb-migration--us-east-1 \ + python scripts/backfill.py + +# 4. Wait for stream replay to drain. Run the gate. +DLQ_URL= python scripts/convergence_check.py + +# 5. Sample-verify before flipping app routing. +python scripts/verify_cutover.py --sample-size 1000 + +# 6. Flip your application's table reference. Resume traffic. + +# 7. After 7-14 days of validation: +python scripts/cleanup.py +./teardown.sh CONFIRM=yes # removes Lambda/role/DLQ/SNS/alarms/bucket +``` + +## Configuration + +`deploy.sh` reads from environment variables or from `./config.env` if present. +A non-exhaustive list: + +| Var | Default | Notes | +|-----|---------|-------| +| `SOURCE_TABLE` | required | Existing table | +| `TARGET_TABLE` | required | Created by deploy.sh unless cross-account | +| `PARTITION_KEY` | `pk` | Source-table partition-key attribute name | +| `PARTITION_KEY_TYPE` | `S` | `S`, `N`, or `B` | +| `SORT_KEY` | (unset) | Leave empty for hash-only schema | +| `SORT_KEY_TYPE` | `S` | | +| `REGION` | `us-east-1` | | +| `LAMBDA_FUNCTION_NAME` | `ddb-migration-stream-replay` | | +| `LAMBDA_ROLE_NAME` | `ddb-migration-stream-replay-role` | | +| `DLQ_NAME` | `ddb-migration-dlq` | | +| `SNS_TOPIC_NAME` | `ddb-migration-alerts` | Subscribe an endpoint after deploy | +| `ITERATOR_AGE_WARN_MS` | `43200000` (12 h) | Warning alarm threshold | +| `ITERATOR_AGE_CRIT_MS` | `72000000` (20 h) | Critical alarm threshold | +| `TARGET_ACCOUNT` | (unset) | Set for cross-account; switches mode | +| `TARGET_ROLE_ARN` | (unset) | Required when `TARGET_ACCOUNT` differs | +| `TRANSFORM_MODULE` | (unset) | Custom Python module path; falls back to bundled `transform.py` | + +`backfill.py` and `convergence_check.py` accept the same env vars plus their +own CLI flags (`--dry-run`, `--ignore-count-drift`, `--max-iterator-age-ms`, +etc.). Run any of them with `--help` for the full list. + +## Customizing the transform + +Edit `transform.py` (or set `TRANSFORM_MODULE` to a different module). The +function runs in the Lambda, in `backfill.py`, and in `verify_cutover.py` — +all three must share the same logic, which is why it is one file. + +```python +def transform(item, source_event=None): + # Rename a column. + if "user_id" in item: + item["customer_id"] = item.pop("user_id") + # Compute a new GSI key. + item["status_idx"] = f"{item['status']}#{item['created_at']}" + return item # or return None to skip the item entirely +``` + +## Cross-account migrations + +When the target table lives in a different account: + +1. In the **target** account, create a role `ddb-migration-target-writer` that + trusts the source-account stream-replay role and grants + `dynamodb:PutItem` + `dynamodb:UpdateItem` on the target table. The + resource-based policy template is in `iam/policies.json` → + `CrossAccountTargetTablePolicy`. +2. In the **source** account, set `TARGET_ACCOUNT` and `TARGET_ROLE_ARN` + before running `deploy.sh`. The Lambda's inline policy will include + `sts:AssumeRole` for that role. +3. To run `backfill.py` from the target account against the source-account + export bucket, attach `iam/policies.json:CrossAccountExportBucketPolicy` to + the bucket. + +Test IAM end-to-end before you start the migration. Permission failures +mid-stream waste the 24-hour stream-retention window. + +## Convergence gates + +`convergence_check.py` runs three checks in sequence: + +1. **Iterator age** is below `--max-iterator-age-ms` (default 1000 ms). +2. **DLQ** is empty (visible + in-flight). +3. **Scan COUNT** on source vs. target is within `--count-drift-pct` + (default 0.5%). The target scan excludes tombstones, so deleted-and-replayed + items don't inflate the count. Use `--ignore-count-drift` to skip this check + on very large tables where the Scan would be expensive. + +The script exits non-zero on any failure, so you can use it in CI: + +```sh +python scripts/convergence_check.py || { echo "not ready"; exit 1; } +``` + +## Demo + +`demo/run_demo.sh` provisions a tiny source/target pair, seeds 10K items, +drives live writes during the migration, and verifies the cutover. End-to-end +runtime ~12 min, cost <$1. See `demo/README.md`. + +## Limitations + +* Streams retain records for 24 h. If the Lambda falls behind by that long, + data is lost — alarms fire at 12 h (warning) and 20 h (critical). +* `Scan COUNT` for very large tables is expensive and time-consuming. Allow + several minutes for tables over 100 GiB. +* Glue-based backfill for tables larger than ~100 GiB is not bundled; + `backfill.py` parallelizes within one host. For very large tables, fan it + out across multiple hosts or write a small Glue wrapper. +* Tombstones live for `--tombstone-ttl-days` (default 7) after `cleanup.py` + runs. They are eventually deleted by DynamoDB TTL. +* Rollback is *not* a flag flip back. Once you cut over, returning to the + source requires deploying a reverse-replay Lambda before the cutover so the + source stays current. + +## Tests + +`make test` runs the suite with moto-backed mocks; no AWS calls. `make +coverage` for a per-line report. The integration path (`demo/run_demo.sh`) is +gated behind `DDB_MIGRATION_DEMO_CONFIRM=yes` and requires real AWS creds. + +## References + +* [AWS docs — Export to S3](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/S3DataExport.html) +* [AWS docs — DynamoDB Streams](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Streams.html) +* [AWS docs — Conditional writes](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.ConditionExpressions.html) +* [Bulk Executor for DynamoDB](https://github.com/awslabs/amazon-dynamodb-tools/tree/main/tools/bulk_executor) — sibling tool for non-zero-downtime bulk operations. diff --git a/tools/ddb_migration/demo/README.md b/tools/ddb_migration/demo/README.md new file mode 100644 index 00000000..742f2e66 --- /dev/null +++ b/tools/ddb_migration/demo/README.md @@ -0,0 +1,59 @@ +# `tools/ddb_migration/demo` — One-click migration demo + +Provisions a small source + target pair, seeds 10,000 items, drives live writes +during the migration, runs the convergence gate, and verifies the cutover. + +End-to-end: ~12–15 minutes. Cost: <$1 in a US region. + +## Prerequisites + +* An AWS account with admin-equivalent permissions (the demo creates IAM roles, + Lambda, DynamoDB tables, S3 bucket, SNS topic, SQS queue, CloudWatch alarms). +* `aws` CLI v2 configured (`AWS_PROFILE` set). +* Python 3.10+ with the parent tool's dependencies installed: + ```sh + cd tools/ddb_migration + make install + source .venv/bin/activate + ``` + +## Run it + +```sh +cp demo/config.example.env demo/config.env +# edit demo/config.env: set AWS_PROFILE and REGION +DDB_MIGRATION_DEMO_CONFIRM=yes ./demo/run_demo.sh +``` + +## What you'll see + +``` +========== 1/9 Creating source table ddb-migration-demo-source ========== +========== 2/9 Provisioning migration infrastructure (deploy.sh) ========== +[deploy] Ensuring target table ddb-migration-demo-target exists in us-east-1 +... +========== 9/9 Verifying sample of items ========== + matched=500 missing=0 diverged=0 total=500 +VERIFY OK + +DEMO PASSED +``` + +## Clean up + +```sh +CONFIRM=yes ./teardown.sh +aws dynamodb delete-table --table-name ddb-migration-demo-source --region us-east-1 +aws dynamodb delete-table --table-name ddb-migration-demo-target --region us-east-1 +``` + +## What it actually demonstrates + +1. Source table receives live writes throughout the migration window. +2. The S3 export captures the source at a point in time. +3. `backfill.py` loads the export into the target with `_migration_ts=0`. +4. The Lambda is replaying live writes (with newer `_migration_ts`) in parallel. +5. Conflict resolution: backfill writes never overwrite live updates. +6. Tombstones prevent the backfill from resurrecting deleted items. +7. The convergence gate blocks cutover until iterator age, DLQ, and counts agree. +8. The verifier samples 500 items and proves source ↔ target parity. diff --git a/tools/ddb_migration/demo/config.example.env b/tools/ddb_migration/demo/config.example.env new file mode 100644 index 00000000..784b23e1 --- /dev/null +++ b/tools/ddb_migration/demo/config.example.env @@ -0,0 +1,15 @@ +# Copy to ./config.env and fill in the blanks before running run_demo.sh. +# Both deploy.sh and run_demo.sh source this file automatically. + +# Required. +export AWS_PROFILE=your-dev-profile +export REGION=us-east-1 + +# Demo-specific overrides — leave defaults unless you want different names. +export SOURCE_TABLE=ddb-migration-demo-source +export TARGET_TABLE=ddb-migration-demo-target +export PARTITION_KEY=pk +export SORT_KEY=sk +export DEMO_ITEM_COUNT=10000 +export DEMO_LIVE_WRITE_RATE=5 +export DEMO_LIVE_WRITE_DURATION_SECS=120 diff --git a/tools/ddb_migration/demo/live_writer.py b/tools/ddb_migration/demo/live_writer.py new file mode 100644 index 00000000..d3adbdbe --- /dev/null +++ b/tools/ddb_migration/demo/live_writer.py @@ -0,0 +1,89 @@ +"""Background writer simulating live application traffic during the demo. + +Writes a configurable rate of mutations to the source table for a fixed +duration. Mix of inserts (new keys), updates (existing keys), and deletes — +exercises every stream event type. Run from run_demo.sh as a background job. +""" + +from __future__ import annotations + +import argparse +import os +import random +import string +import sys +import time +from decimal import Decimal + +import boto3 +from botocore.exceptions import ClientError + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--table", default=os.environ.get("SOURCE_TABLE")) + p.add_argument("--partition-key", default=os.environ.get("PARTITION_KEY", "pk")) + p.add_argument("--sort-key", default=os.environ.get("SORT_KEY") or None) + p.add_argument("--rate", type=int, default=int(os.environ.get("DEMO_LIVE_WRITE_RATE", "5"))) + p.add_argument("--duration", type=int, default=int(os.environ.get("DEMO_LIVE_WRITE_DURATION_SECS", "120"))) + p.add_argument("--seed-count", type=int, default=int(os.environ.get("DEMO_ITEM_COUNT", "10000"))) + args = p.parse_args() + if not args.table: + print("SOURCE_TABLE is required", file=sys.stderr) + return 2 + + table = boto3.resource("dynamodb", region_name=args.region).Table(args.table) + deadline = time.time() + args.duration + interval = 1.0 / args.rate + counts = {"insert": 0, "update": 0, "delete": 0, "errors": 0} + cycle = 0 + + while time.time() < deadline: + op = random.choices(["insert", "update", "delete"], weights=[2, 6, 2])[0] + try: + if op == "insert": + pk = f"order#new#{int(time.time()*1000)}#{cycle:05d}" + item: dict = { + args.partition_key: pk, + "customer_id": f"cust#{random.randint(1, 1000):04d}", + "amount": Decimal(str(round(random.uniform(1.0, 999.99), 2))), + "status": "NEW", + "notes": "".join(random.choices(string.ascii_lowercase, k=20)), + } + if args.sort_key: + item[args.sort_key] = f"line#{random.randint(1, 5):03d}" + table.put_item(Item=item) + counts["insert"] += 1 + elif op == "update": + idx = random.randint(0, args.seed_count - 1) + key: dict = {args.partition_key: f"order#{idx:08d}"} + if args.sort_key: + key[args.sort_key] = f"line#{random.randint(1, 5):03d}" + table.update_item( + Key=key, + UpdateExpression="SET #s = :s", + ExpressionAttributeNames={"#s": "status"}, + ExpressionAttributeValues={":s": random.choice(["PAID", "SHIPPED", "DELIVERED"])}, + ) + counts["update"] += 1 + else: + idx = random.randint(0, args.seed_count - 1) + key = {args.partition_key: f"order#{idx:08d}"} + if args.sort_key: + key[args.sort_key] = f"line#{random.randint(1, 5):03d}" + table.delete_item(Key=key) + counts["delete"] += 1 + except ClientError as e: + counts["errors"] += 1 + if e.response["Error"]["Code"] not in ("ConditionalCheckFailedException", "ResourceNotFoundException"): + print(f" live_writer error: {e.response['Error']['Code']}", file=sys.stderr) + cycle += 1 + time.sleep(interval) + + print(f"live_writer done: insert={counts['insert']} update={counts['update']} delete={counts['delete']} errors={counts['errors']}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/demo/run_demo.sh b/tools/ddb_migration/demo/run_demo.sh new file mode 100755 index 00000000..90fccb7f --- /dev/null +++ b/tools/ddb_migration/demo/run_demo.sh @@ -0,0 +1,133 @@ +#!/usr/bin/env bash +# One-click demo of the zero-downtime migration toolkit. +# +# Provisions a tiny source/target pair, seeds 10K items, drives live writes +# during the migration, runs the convergence gate, then verifies the cutover. +# End-to-end runtime: ~12-15 minutes. Estimated cost: <$1 in a US region. +# +# Requires DDB_MIGRATION_DEMO_CONFIRM=yes — this provisions REAL AWS resources. +# Run ../teardown.sh CONFIRM=yes when done. + +set -euo pipefail + +DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TOOL_DIR="$(cd "$DEMO_DIR/.." && pwd)" + +if [[ -f "$DEMO_DIR/config.env" ]]; then + # shellcheck source=/dev/null + source "$DEMO_DIR/config.env" +elif [[ -f "$TOOL_DIR/config.env" ]]; then + # shellcheck source=/dev/null + source "$TOOL_DIR/config.env" +else + echo "ERROR: copy demo/config.example.env to demo/config.env first" >&2 + exit 2 +fi + +if [[ "${DDB_MIGRATION_DEMO_CONFIRM:-}" != "yes" ]]; then + cat >&2 </dev/null 2>&1; then + # shellcheck disable=SC2086 + aws dynamodb create-table \ + --table-name "$SOURCE_TABLE" \ + --attribute-definitions $attr_defs \ + --key-schema $key_schema \ + --billing-mode PAY_PER_REQUEST \ + --region "$REGION" >/dev/null + aws dynamodb wait table-exists --table-name "$SOURCE_TABLE" --region "$REGION" +fi +aws dynamodb update-continuous-backups \ + --table-name "$SOURCE_TABLE" \ + --point-in-time-recovery-specification PointInTimeRecoveryEnabled=true \ + --region "$REGION" >/dev/null 2>&1 || true + +step "2/9 Provisioning migration infrastructure (deploy.sh)" +( cd "$TOOL_DIR" && bash deploy.sh ) + +step "3/9 Seeding ${DEMO_ITEM_COUNT:-10000} items into $SOURCE_TABLE" +python "$DEMO_DIR/seed_data.py" + +step "4/9 Triggering S3 export" +ACCOUNT="$(aws sts get-caller-identity --query Account --output text)" +EXPORT_BUCKET="${EXPORT_BUCKET:-ddb-migration-${ACCOUNT}-${REGION}}" +SOURCE_TABLE_ARN="$(aws dynamodb describe-table --table-name "$SOURCE_TABLE" --region "$REGION" --query 'Table.TableArn' --output text)" +EXPORT_ARN="$(aws dynamodb export-table-to-point-in-time \ + --table-arn "$SOURCE_TABLE_ARN" \ + --s3-bucket "$EXPORT_BUCKET" \ + --s3-prefix exports/ \ + --export-format DYNAMODB_JSON \ + --region "$REGION" \ + --query 'ExportDescription.ExportArn' --output text)" +echo " export ARN: $EXPORT_ARN" + +step "5/9 Starting live-writer in background (rate=${DEMO_LIVE_WRITE_RATE:-5}/s, duration=${DEMO_LIVE_WRITE_DURATION_SECS:-120}s)" +python "$DEMO_DIR/live_writer.py" & +LIVE_WRITER_PID=$! +trap 'kill $LIVE_WRITER_PID 2>/dev/null || true' EXIT + +step "6/9 Waiting for export to complete" +while true; do + status="$(aws dynamodb describe-export --export-arn "$EXPORT_ARN" --region "$REGION" --query 'ExportDescription.ExportStatus' --output text)" + echo " export status: $status" + [[ "$status" == "COMPLETED" ]] && break + [[ "$status" == "FAILED" ]] && { echo " export failed"; exit 1; } + sleep 30 +done + +step "7/9 Running backfill" +EXPORT_BUCKET="$EXPORT_BUCKET" python "$TOOL_DIR/scripts/backfill.py" + +step "8/9 Waiting for live-writer to finish, then running convergence check" +wait "$LIVE_WRITER_PID" || true +trap - EXIT +DLQ_URL="$(aws sqs get-queue-url --queue-name ddb-migration-dlq --region "$REGION" --query 'QueueUrl' --output text)" +# 240s wait with 120s idle grace covers the case where the demo's traffic has stopped +# and the Lambda is no longer emitting IteratorAge datapoints. +DLQ_URL="$DLQ_URL" python "$TOOL_DIR/scripts/convergence_check.py" --max-wait-seconds 240 + +step "9/9 Verifying sample of items" +python "$TOOL_DIR/scripts/verify_cutover.py" --sample-size 500 + +cat < dict: + item: dict = { + partition_key: f"order#{idx:08d}", + "customer_id": f"cust#{random.randint(1, 1000):04d}", + "amount": Decimal(str(round(random.uniform(1.0, 999.99), 2))), + "status": random.choice(["NEW", "PAID", "SHIPPED", "DELIVERED"]), + "notes": "".join(random.choices(string.ascii_lowercase, k=20)), + } + if sort_key: + item[sort_key] = f"line#{random.randint(1, 5):03d}" + return item + + +def write_batch(table, items: list[dict]) -> None: + with table.batch_writer() as batch: + for item in items: + batch.put_item(Item=item) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--table", default=os.environ.get("SOURCE_TABLE")) + p.add_argument("--partition-key", default=os.environ.get("PARTITION_KEY", "pk")) + p.add_argument("--sort-key", default=os.environ.get("SORT_KEY") or None) + p.add_argument("--count", type=int, default=int(os.environ.get("DEMO_ITEM_COUNT", "10000"))) + p.add_argument("--workers", type=int, default=8) + args = p.parse_args() + if not args.table: + print("SOURCE_TABLE is required", file=sys.stderr) + return 2 + table = boto3.resource("dynamodb", region_name=args.region).Table(args.table) + print(f"Seeding {args.count} items into {args.table}...") + chunk = max(1, args.count // args.workers) + batches: list[list[dict]] = [] + cur: list[dict] = [] + for i in range(args.count): + cur.append(make_item(i, args.partition_key, args.sort_key)) + if len(cur) >= 25: + batches.append(cur) + cur = [] + if cur: + batches.append(cur) + with ThreadPoolExecutor(max_workers=args.workers) as pool: + for _ in pool.map(lambda b: write_batch(table, b), batches): + pass + print(f"Seeded {args.count} items") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/deploy.sh b/tools/ddb_migration/deploy.sh new file mode 100755 index 00000000..78051084 --- /dev/null +++ b/tools/ddb_migration/deploy.sh @@ -0,0 +1,348 @@ +#!/usr/bin/env bash +# Provisions the DynamoDB zero-downtime migration infrastructure. +# +# Idempotent: safe to re-run. Reads ./config.env if present, otherwise relies on +# environment variables. All AWS API calls use the AWS_PROFILE / AWS_REGION +# already in your shell. +# +# Required env vars: +# SOURCE_TABLE Existing source table. +# TARGET_TABLE Target table to create (or already exists). +# +# Optional env vars (with defaults): +# PARTITION_KEY=pk +# PARTITION_KEY_TYPE=S # S | N | B +# SORT_KEY= # leave empty for hash-only schema +# SORT_KEY_TYPE=S +# REGION=us-east-1 +# EXPORT_BUCKET=ddb-migration-- +# LAMBDA_FUNCTION_NAME=ddb-migration-stream-replay +# LAMBDA_ROLE_NAME=ddb-migration-stream-replay-role +# DLQ_NAME=ddb-migration-dlq +# SNS_TOPIC_NAME=ddb-migration-alerts +# ITERATOR_AGE_WARN_MS=43200000 # 12h +# ITERATOR_AGE_CRIT_MS=72000000 # 20h +# +# Cross-account (target table in a different account): +# TARGET_ACCOUNT=123456789012 # account that owns TARGET_TABLE +# TARGET_ROLE_ARN=arn:aws:iam::123456789012:role/ddb-migration-target-writer +# # role the Lambda assumes to write to target +# # (must be created in the target account first) +# +# After running this script, the Lambda is wired to the source-table stream and +# replay starts immediately. Subscribe to the printed SNS topic ARN to receive +# IteratorAge alarm notifications. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -f "$SCRIPT_DIR/config.env" ]]; then + # shellcheck source=/dev/null + source "$SCRIPT_DIR/config.env" +fi + +: "${SOURCE_TABLE:?SOURCE_TABLE is required}" +: "${TARGET_TABLE:?TARGET_TABLE is required}" + +PARTITION_KEY="${PARTITION_KEY:-pk}" +PARTITION_KEY_TYPE="${PARTITION_KEY_TYPE:-S}" +SORT_KEY="${SORT_KEY:-}" +SORT_KEY_TYPE="${SORT_KEY_TYPE:-S}" +REGION="${REGION:-us-east-1}" +LAMBDA_FUNCTION_NAME="${LAMBDA_FUNCTION_NAME:-ddb-migration-stream-replay}" +LAMBDA_ROLE_NAME="${LAMBDA_ROLE_NAME:-ddb-migration-stream-replay-role}" +DLQ_NAME="${DLQ_NAME:-ddb-migration-dlq}" +SNS_TOPIC_NAME="${SNS_TOPIC_NAME:-ddb-migration-alerts}" +ITERATOR_AGE_WARN_MS="${ITERATOR_AGE_WARN_MS:-43200000}" +ITERATOR_AGE_CRIT_MS="${ITERATOR_AGE_CRIT_MS:-72000000}" + +ACCOUNT="$(aws sts get-caller-identity --query Account --output text)" +EXPORT_BUCKET="${EXPORT_BUCKET:-ddb-migration-${ACCOUNT}-${REGION}}" + +CROSS_ACCOUNT="false" +if [[ -n "${TARGET_ACCOUNT:-}" && "$TARGET_ACCOUNT" != "$ACCOUNT" ]]; then + CROSS_ACCOUNT="true" + : "${TARGET_ROLE_ARN:?TARGET_ROLE_ARN is required for cross-account migrations}" +fi + +log() { printf "\n[deploy] %s\n" "$*"; } + +# +# 1. Target table +# +if [[ "$CROSS_ACCOUNT" == "true" ]]; then + log "Cross-account mode: skipping target table creation (must exist in account $TARGET_ACCOUNT)" +else + log "Ensuring target table $TARGET_TABLE exists in $REGION" + if aws dynamodb describe-table --table-name "$TARGET_TABLE" --region "$REGION" >/dev/null 2>&1; then + log " already exists" + else + attr_defs="AttributeName=$PARTITION_KEY,AttributeType=$PARTITION_KEY_TYPE" + key_schema="AttributeName=$PARTITION_KEY,KeyType=HASH" + if [[ -n "$SORT_KEY" ]]; then + attr_defs="$attr_defs AttributeName=$SORT_KEY,AttributeType=$SORT_KEY_TYPE" + key_schema="$key_schema AttributeName=$SORT_KEY,KeyType=RANGE" + fi + # shellcheck disable=SC2086 + aws dynamodb create-table \ + --table-name "$TARGET_TABLE" \ + --attribute-definitions $attr_defs \ + --key-schema $key_schema \ + --billing-mode PAY_PER_REQUEST \ + --region "$REGION" >/dev/null + aws dynamodb wait table-exists --table-name "$TARGET_TABLE" --region "$REGION" + log " created" + fi + log "Enabling TTL on target table (_ttl attribute) for tombstone expiration" + aws dynamodb update-time-to-live \ + --table-name "$TARGET_TABLE" \ + --time-to-live-specification "Enabled=true,AttributeName=_ttl" \ + --region "$REGION" >/dev/null 2>&1 || true +fi + +# +# 2. Streams on source +# +log "Enabling DynamoDB Streams on source table $SOURCE_TABLE" +aws dynamodb update-table \ + --table-name "$SOURCE_TABLE" \ + --stream-specification StreamEnabled=true,StreamViewType=NEW_AND_OLD_IMAGES \ + --region "$REGION" >/dev/null 2>&1 || log " (already enabled or other error)" +SOURCE_STREAM_ARN="$(aws dynamodb describe-table --table-name "$SOURCE_TABLE" --region "$REGION" --query 'Table.LatestStreamArn' --output text)" +log " stream ARN: $SOURCE_STREAM_ARN" + +# +# 3. S3 export bucket +# +log "Ensuring export bucket s3://$EXPORT_BUCKET exists" +if aws s3api head-bucket --bucket "$EXPORT_BUCKET" --region "$REGION" 2>/dev/null; then + log " already exists" +else + if [[ "$REGION" == "us-east-1" ]]; then + aws s3api create-bucket --bucket "$EXPORT_BUCKET" --region "$REGION" >/dev/null + else + aws s3api create-bucket \ + --bucket "$EXPORT_BUCKET" \ + --region "$REGION" \ + --create-bucket-configuration "LocationConstraint=$REGION" >/dev/null + fi + aws s3api put-bucket-encryption \ + --bucket "$EXPORT_BUCKET" \ + --server-side-encryption-configuration '{"Rules":[{"ApplyServerSideEncryptionByDefault":{"SSEAlgorithm":"AES256"}}]}' >/dev/null + log " created" +fi + +# +# 4. SNS topic for alarm actions +# +log "Ensuring SNS topic $SNS_TOPIC_NAME exists" +SNS_TOPIC_ARN="$(aws sns create-topic --name "$SNS_TOPIC_NAME" --region "$REGION" --query 'TopicArn' --output text)" +log " topic ARN: $SNS_TOPIC_ARN" + +# +# 5. SQS DLQ for stream-replay failures +# +log "Ensuring SQS DLQ $DLQ_NAME exists" +DLQ_URL="$(aws sqs create-queue --queue-name "$DLQ_NAME" --region "$REGION" --query 'QueueUrl' --output text)" +DLQ_ARN="$(aws sqs get-queue-attributes --queue-url "$DLQ_URL" --attribute-names QueueArn --region "$REGION" --query 'Attributes.QueueArn' --output text)" +log " DLQ URL: $DLQ_URL" +log " DLQ ARN: $DLQ_ARN" + +# +# 6. IAM role for the Lambda +# +log "Ensuring IAM role $LAMBDA_ROLE_NAME exists" +trust_policy=$(cat </dev/null 2>&1; then + log " already exists" +else + aws iam create-role \ + --role-name "$LAMBDA_ROLE_NAME" \ + --assume-role-policy-document "$trust_policy" >/dev/null + log " created; waiting for IAM propagation..." + sleep 10 +fi + +inline_policy=$(cat </dev/null +ROLE_ARN="$(aws iam get-role --role-name "$LAMBDA_ROLE_NAME" --query 'Role.Arn' --output text)" +log " role ARN: $ROLE_ARN" + +# +# 7. Package + deploy Lambda +# +log "Packaging Lambda" +PKG_DIR="$(mktemp -d)" +cp "$SCRIPT_DIR/lambda/stream_replay.py" "$PKG_DIR/" +cp "$SCRIPT_DIR/transform.py" "$PKG_DIR/" +( cd "$PKG_DIR" && zip -q -r "$SCRIPT_DIR/lambda.zip" . ) +rm -rf "$PKG_DIR" + +env_vars="Variables={TARGET_TABLE=$TARGET_TABLE,PARTITION_KEY=$PARTITION_KEY,TARGET_REGION=$REGION" +if [[ "$CROSS_ACCOUNT" == "true" ]]; then + env_vars="${env_vars},TARGET_ROLE_ARN=$TARGET_ROLE_ARN" +fi +env_vars="${env_vars}}" + +if aws lambda get-function --function-name "$LAMBDA_FUNCTION_NAME" --region "$REGION" >/dev/null 2>&1; then + log "Updating Lambda code" + aws lambda update-function-code \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --zip-file "fileb://$SCRIPT_DIR/lambda.zip" \ + --region "$REGION" >/dev/null + aws lambda wait function-updated --function-name "$LAMBDA_FUNCTION_NAME" --region "$REGION" + aws lambda update-function-configuration \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --environment "$env_vars" \ + --region "$REGION" >/dev/null +else + log "Creating Lambda $LAMBDA_FUNCTION_NAME" + aws lambda create-function \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --runtime python3.12 \ + --role "$ROLE_ARN" \ + --handler stream_replay.handler \ + --timeout 300 \ + --memory-size 512 \ + --zip-file "fileb://$SCRIPT_DIR/lambda.zip" \ + --environment "$env_vars" \ + --region "$REGION" >/dev/null + aws lambda wait function-active --function-name "$LAMBDA_FUNCTION_NAME" --region "$REGION" +fi +rm -f "$SCRIPT_DIR/lambda.zip" + +# +# 8. Event source mapping (with DLQ) +# +log "Ensuring event source mapping is in place" +EXISTING_ESM="$(aws lambda list-event-source-mappings \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --region "$REGION" \ + --query "EventSourceMappings[?EventSourceArn=='$SOURCE_STREAM_ARN'].UUID | [0]" \ + --output text)" +if [[ "$EXISTING_ESM" == "None" || -z "$EXISTING_ESM" ]]; then + aws lambda create-event-source-mapping \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --event-source-arn "$SOURCE_STREAM_ARN" \ + --starting-position TRIM_HORIZON \ + --batch-size 100 \ + --maximum-batching-window-in-seconds 5 \ + --bisect-batch-on-function-error \ + --maximum-retry-attempts 3 \ + --function-response-types ReportBatchItemFailures \ + --destination-config "OnFailure={Destination=$DLQ_ARN}" \ + --region "$REGION" >/dev/null + log " created" +else + log " already exists ($EXISTING_ESM); updating destination config" + aws lambda update-event-source-mapping \ + --uuid "$EXISTING_ESM" \ + --destination-config "OnFailure={Destination=$DLQ_ARN}" \ + --region "$REGION" >/dev/null +fi + +# +# 9. CloudWatch alarms wired to SNS +# +log "Provisioning IteratorAge alarms (warn=$ITERATOR_AGE_WARN_MS ms, crit=$ITERATOR_AGE_CRIT_MS ms)" +for tier in WARN CRIT; do + if [[ "$tier" == "WARN" ]]; then + threshold="$ITERATOR_AGE_WARN_MS" + name="ddb-migration-iterator-age-warning" + else + threshold="$ITERATOR_AGE_CRIT_MS" + name="ddb-migration-iterator-age-critical" + fi + aws cloudwatch put-metric-alarm \ + --alarm-name "$name" \ + --alarm-description "Stream-replay Lambda IteratorAge (tier $tier)" \ + --namespace "AWS/Lambda" \ + --metric-name IteratorAge \ + --dimensions "Name=FunctionName,Value=$LAMBDA_FUNCTION_NAME" \ + --statistic Maximum \ + --period 60 \ + --evaluation-periods 5 \ + --threshold "$threshold" \ + --comparison-operator GreaterThanThreshold \ + --alarm-actions "$SNS_TOPIC_ARN" \ + --treat-missing-data notBreaching \ + --region "$REGION" >/dev/null +done +log " done" + +cat < None: + payload = {"ts": time.time(), "level": level, "event": event, **fields} + logger.info(json.dumps(payload, default=str)) + + +def _load_transform() -> Callable[..., Any]: + """Resolve the user-supplied transform function. + + Order: ``TRANSFORM_MODULE`` env var, then bundled ``transform`` module, then + identity fallback. + """ + module_name = os.environ.get("TRANSFORM_MODULE") + if module_name: + try: + mod = importlib.import_module(module_name) + fn = getattr(mod, "transform") + _log("info", "transform_loaded", source=module_name) + return fn + except (ImportError, AttributeError) as e: + _log("warning", "transform_load_failed", source=module_name, error=str(e)) + try: + # Allow co-located transform.py when packaged with the Lambda zip. + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + from transform import transform as fn # type: ignore + + _log("info", "transform_loaded", source="bundled") + return fn + except ImportError: + _log("info", "transform_loaded", source="identity") + return lambda item, source_event=None: item + + +# Module-level setup so cold-start cost is paid once. +_TARGET_TABLE = os.environ.get("TARGET_TABLE") +_PARTITION_KEY = os.environ.get("PARTITION_KEY", "pk") +_TARGET_REGION = os.environ.get("TARGET_REGION") or os.environ.get("AWS_REGION") or "us-east-1" +_TARGET_ROLE_ARN = os.environ.get("TARGET_ROLE_ARN") +_DESERIALIZER = TypeDeserializer() +_TRANSFORM = _load_transform() + + +def _build_target_client() -> Any: + """Build a DynamoDB resource client, optionally assuming a cross-account role.""" + if _TARGET_ROLE_ARN: + sts = boto3.client("sts") + creds = sts.assume_role( + RoleArn=_TARGET_ROLE_ARN, + RoleSessionName="ddb-migration-stream-replay", + DurationSeconds=3600, + )["Credentials"] + return boto3.resource( + "dynamodb", + region_name=_TARGET_REGION, + aws_access_key_id=creds["AccessKeyId"], + aws_secret_access_key=creds["SecretAccessKey"], + aws_session_token=creds["SessionToken"], + ) + return boto3.resource("dynamodb", region_name=_TARGET_REGION) + + +_TABLE = _build_target_client().Table(_TARGET_TABLE) if _TARGET_TABLE else None + + +def _deserialize_image(image: dict[str, Any]) -> dict[str, Any]: + return {k: _DESERIALIZER.deserialize(v) for k, v in image.items()} + + +def _conditional_put(item: dict[str, Any], migration_ts: float) -> str: + """Put item with newer-wins condition. Returns 'written' or 'skipped'. + + Migration timestamps are stored as Decimal because the boto3 resource API + rejects native float values. + """ + ts = Decimal(str(migration_ts)) + item = dict(item) + item["_migration_ts"] = ts + try: + _TABLE.put_item( + Item=item, + ConditionExpression="attribute_not_exists(#pk) OR #ts < :ts", + ExpressionAttributeNames={"#pk": _PARTITION_KEY, "#ts": "_migration_ts"}, + ExpressionAttributeValues={":ts": ts}, + ) + return "written" + except ClientError as e: + if e.response["Error"]["Code"] == "ConditionalCheckFailedException": + return "skipped" + raise + + +def handler(event: dict[str, Any], context: Any) -> dict[str, Any]: + """Lambda entry point. Returns ``batchItemFailures`` for partial-batch retry.""" + failures: list[dict[str, str]] = [] + counts = {"written": 0, "skipped": 0, "tombstone": 0, "errors": 0} + + for record in event.get("Records", []): + event_id = record.get("eventID", "") + event_name = record.get("eventName", "") + ddb = record.get("dynamodb", {}) + migration_ts = float(ddb.get("ApproximateCreationDateTime", time.time())) + try: + if event_name in ("INSERT", "MODIFY"): + new_image = ddb.get("NewImage") + if not new_image: + continue + item = _deserialize_image(new_image) + transformed = _TRANSFORM(item, record) + if transformed is None: + continue + outcome = _conditional_put(transformed, migration_ts) + counts[outcome] += 1 + elif event_name == "REMOVE": + old_image = ddb.get("OldImage") + if not old_image: + continue + key_only = _deserialize_image(old_image) + tombstone = {_PARTITION_KEY: key_only[_PARTITION_KEY], "_tombstone": True} + # Preserve sort key if present; deserialize_image already did it. + for k, v in key_only.items(): + if k == _PARTITION_KEY or k.startswith("_"): + continue + # Sort key heuristic: any other top-level key from the item's primary key. + if "Keys" in ddb and k in _deserialize_image(ddb["Keys"]): + tombstone[k] = v + _conditional_put(tombstone, migration_ts) + counts["tombstone"] += 1 + except Exception as e: # noqa: BLE001 — partial-batch retry needs broad catch + counts["errors"] += 1 + _log("error", "record_failed", event_id=event_id, event_name=event_name, error=str(e)) + failures.append({"itemIdentifier": event_id}) + + _log("info", "batch_complete", **counts, failures=len(failures)) + return {"batchItemFailures": failures} diff --git a/tools/ddb_migration/pytest.ini b/tools/ddb_migration/pytest.ini new file mode 100644 index 00000000..ccba29c7 --- /dev/null +++ b/tools/ddb_migration/pytest.ini @@ -0,0 +1,6 @@ +[pytest] +pythonpath = . scripts lambda +testpaths = tests +python_files = test_*.py +python_classes = Test* +python_functions = test_* diff --git a/tools/ddb_migration/requirements-dev.txt b/tools/ddb_migration/requirements-dev.txt new file mode 100644 index 00000000..d5fcb4ef --- /dev/null +++ b/tools/ddb_migration/requirements-dev.txt @@ -0,0 +1,4 @@ +-r requirements.txt +pytest>=7.4.0 +pytest-cov>=4.1.0 +moto[dynamodb,s3,sqs,cloudwatch]>=4.2.0 diff --git a/tools/ddb_migration/requirements.txt b/tools/ddb_migration/requirements.txt new file mode 100644 index 00000000..a028272f --- /dev/null +++ b/tools/ddb_migration/requirements.txt @@ -0,0 +1 @@ +boto3>=1.33.13 diff --git a/tools/ddb_migration/scripts/backfill.py b/tools/ddb_migration/scripts/backfill.py new file mode 100644 index 00000000..e84f41dd --- /dev/null +++ b/tools/ddb_migration/scripts/backfill.py @@ -0,0 +1,279 @@ +"""Standalone backfill from a DynamoDB S3 export to the target table. + +Use this for tables up to ~100 GiB. For larger tables, fan out across multiple +hosts or use a Glue job (not included in this toolkit). + +Reads the export manifest under ``s3://$EXPORT_BUCKET/$EXPORT_PREFIX/``, +deserializes each gzipped JSONL data file, applies ``transform()``, and writes +to the target table with ``_migration_ts=0`` so any stream-replayed item beats +the backfill version. + +Includes a circuit breaker that pauses the backfill when the stream-replay +Lambda's ``IteratorAge`` rises above 18h (default), preventing the backfill +from consuming all the target table's write capacity and starving the Lambda. + +Configuration via environment variables (see Quick Start in README) or CLI +flags. CLI flags take precedence. +""" + +from __future__ import annotations + +import argparse +import gzip +import importlib +import io +import json +import logging +import os +import random +import sys +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime, timedelta, timezone +from typing import Any, Callable + +import boto3 +from boto3.dynamodb.types import TypeDeserializer +from botocore.exceptions import ClientError + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger("backfill") + + +def load_transform() -> Callable[..., Any]: + """Resolve transform: TRANSFORM_MODULE env var → bundled transform.py → identity.""" + module_name = os.environ.get("TRANSFORM_MODULE") + if module_name: + try: + mod = importlib.import_module(module_name) + return getattr(mod, "transform") + except (ImportError, AttributeError) as e: + log.warning("transform module %s failed to load: %s", module_name, e) + try: + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + from transform import transform # type: ignore + + return transform + except ImportError: + return lambda item, source_event=None: item + + +def get_data_file_keys(s3_client: Any, bucket: str, prefix: str) -> list[str]: + """Discover .json.gz data files via the most recent export's manifest. + + DynamoDB exports lay down keys as + ``{prefix}/AWSDynamoDB/{exportId}/manifest-summary.json``. We list all + ``manifest-summary.json`` keys under the prefix and pick the lexically + largest (export IDs include a timestamp prefix so this picks the newest). + """ + prefix = prefix.rstrip("/") + "/" + paginator = s3_client.get_paginator("list_objects_v2") + summaries: list[str] = [] + for page in paginator.paginate(Bucket=bucket, Prefix=prefix): + for obj in page.get("Contents", []) or []: + if obj["Key"].endswith("/manifest-summary.json"): + summaries.append(obj["Key"]) + if not summaries: + raise RuntimeError(f"No manifest-summary.json found under s3://{bucket}/{prefix}") + summary_key = sorted(summaries)[-1] + summary_obj = s3_client.get_object(Bucket=bucket, Key=summary_key) + summary = json.loads(summary_obj["Body"].read()) + manifest_key = summary["manifestFilesS3Key"] + manifest_obj = s3_client.get_object(Bucket=bucket, Key=manifest_key) + keys: list[str] = [] + for line in manifest_obj["Body"].read().decode().splitlines(): + if line.strip(): + keys.append(json.loads(line)["dataFileS3Key"]) + log.info("Discovered %d data files via %s", len(keys), summary_key) + return keys + + +def deserialize_item(dynamo_item: dict[str, Any], deserializer: TypeDeserializer) -> dict[str, Any]: + return {k: deserializer.deserialize(v) for k, v in dynamo_item.items()} + + +def put_item_with_retry( + table: Any, + item: dict[str, Any], + partition_key: str, + max_retries: int = 8, +) -> str: + """Conditional newer-wins put with exponential backoff. Returns 'written' or 'skipped'.""" + item = dict(item) + item["_migration_ts"] = 0 + for attempt in range(max_retries): + try: + table.put_item( + Item=item, + ConditionExpression="attribute_not_exists(#pk) OR #ts < :ts", + ExpressionAttributeNames={"#pk": partition_key, "#ts": "_migration_ts"}, + ExpressionAttributeValues={":ts": 0}, + ) + return "written" + except ClientError as e: + code = e.response["Error"]["Code"] + if code == "ConditionalCheckFailedException": + return "skipped" + if code in ("ProvisionedThroughputExceededException", "ThrottlingException"): + time.sleep(min(2**attempt * 0.1, 30) + random.uniform(0, 0.5)) + continue + raise + raise RuntimeError(f"Exhausted {max_retries} retries for item") + + +def should_pause( + cw_client: Any, + lambda_function_name: str, + pause_threshold_hours: float, +) -> bool: + """Query CloudWatch for stream-replay IteratorAge; pause if above threshold.""" + end = datetime.now(timezone.utc) + start = end - timedelta(minutes=5) + try: + resp = cw_client.get_metric_statistics( + Namespace="AWS/Lambda", + MetricName="IteratorAge", + Dimensions=[{"Name": "FunctionName", "Value": lambda_function_name}], + StartTime=start, + EndTime=end, + Period=60, + Statistics=["Maximum"], + ) + except ClientError as e: + log.warning("CloudWatch query failed (%s); not pausing", e.response["Error"]["Code"]) + return False + points = resp.get("Datapoints", []) + if not points: + return False + max_age_ms = max(p["Maximum"] for p in points) + threshold_ms = pause_threshold_hours * 3600 * 1000 + if max_age_ms > threshold_ms: + log.warning("IteratorAge %.1fh > threshold %.1fh; pausing", max_age_ms / 3.6e6, pause_threshold_hours) + return True + return False + + +def process_data_file( + s3_client: Any, + cw_client: Any, + table: Any, + bucket: str, + s3_key: str, + config: dict[str, Any], +) -> dict[str, int]: + """Stream a single .json.gz export file into the target.""" + deserializer = TypeDeserializer() + transform_fn = config["transform"] + counts = {"items": 0, "written": 0, "skipped": 0, "errors": 0} + obj = s3_client.get_object(Bucket=bucket, Key=s3_key) + raw = obj["Body"].read() + decompressed = gzip.GzipFile(fileobj=io.BytesIO(raw)).read().decode() + items: list[dict[str, Any]] = [] + for line in decompressed.splitlines(): + if not line.strip(): + continue + items.append(json.loads(line)["Item"]) + # Shuffle to spread writes across partitions and avoid hot-shard throttling. + random.shuffle(items) + if config["dry_run"]: + log.info("[dry-run] %s: %d items", s3_key, len(items)) + counts["items"] = len(items) + return counts + cadence = config["circuit_breaker_check_interval"] + for idx, raw_item in enumerate(items, 1): + if idx % cadence == 0 and should_pause(cw_client, config["lambda_function_name"], config["pause_threshold_hours"]): + while should_pause(cw_client, config["lambda_function_name"], config["pause_threshold_hours"]): + time.sleep(60) + item = deserialize_item(raw_item, deserializer) + transformed = transform_fn(item) + if transformed is None: + counts["items"] += 1 + continue + try: + outcome = put_item_with_retry(table, transformed, config["partition_key"]) + counts[outcome] += 1 + except Exception as e: # noqa: BLE001 + counts["errors"] += 1 + log.exception("Failed to write item from %s: %s", s3_key, e) + counts["items"] += 1 + log.info( + "Done %s: %d items, %d written, %d skipped, %d errors", + s3_key, counts["items"], counts["written"], counts["skipped"], counts["errors"], + ) + return counts + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument("--export-bucket", default=os.environ.get("EXPORT_BUCKET")) + p.add_argument("--export-prefix", default=os.environ.get("EXPORT_PREFIX", "exports/")) + p.add_argument("--target-table", default=os.environ.get("TARGET_TABLE")) + p.add_argument("--partition-key", default=os.environ.get("PARTITION_KEY", "pk")) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--max-workers", type=int, default=int(os.environ.get("MAX_WORKERS", "16"))) + p.add_argument( + "--lambda-function-name", + default=os.environ.get("LAMBDA_FUNCTION_NAME", "ddb-migration-stream-replay"), + ) + p.add_argument( + "--pause-threshold-hours", + type=float, + default=float(os.environ.get("ITERATOR_AGE_PAUSE_HOURS", "18")), + ) + p.add_argument( + "--circuit-breaker-check-interval", + type=int, + default=int(os.environ.get("CIRCUIT_BREAKER_CHECK_INTERVAL", "5000")), + ) + p.add_argument("--dry-run", action="store_true", help="Parse files, count items, write nothing") + return p.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if not args.export_bucket or not args.target_table: + log.error("EXPORT_BUCKET and TARGET_TABLE are required") + return 2 + s3_client = boto3.client("s3", region_name=args.region) + cw_client = boto3.client("cloudwatch", region_name=args.region) + table = boto3.resource("dynamodb", region_name=args.region).Table(args.target_table) + config = { + "transform": load_transform(), + "partition_key": args.partition_key, + "lambda_function_name": args.lambda_function_name, + "pause_threshold_hours": args.pause_threshold_hours, + "circuit_breaker_check_interval": args.circuit_breaker_check_interval, + "dry_run": args.dry_run, + } + keys = get_data_file_keys(s3_client, args.export_bucket, args.export_prefix) + if not keys: + log.warning("No data files found; nothing to backfill") + return 0 + totals = {"items": 0, "written": 0, "skipped": 0, "errors": 0, "files_failed": 0} + start = time.time() + with ThreadPoolExecutor(max_workers=args.max_workers) as pool: + futures = { + pool.submit(process_data_file, s3_client, cw_client, table, args.export_bucket, key, config): key + for key in keys + } + for fut in as_completed(futures): + key = futures[fut] + try: + counts = fut.result() + for k, v in counts.items(): + totals[k] += v + except Exception as e: # noqa: BLE001 + totals["files_failed"] += 1 + log.exception("File failed %s: %s", key, e) + elapsed = time.time() - start + rate = totals["items"] / elapsed if elapsed > 0 else 0 + log.info( + "BACKFILL COMPLETE — items=%d written=%d skipped=%d errors=%d files_failed=%d elapsed=%.1fs rate=%.0f/s", + totals["items"], totals["written"], totals["skipped"], totals["errors"], + totals["files_failed"], elapsed, rate, + ) + return 1 if totals["files_failed"] > 0 or totals["errors"] > 0 else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/scripts/cleanup.py b/tools/ddb_migration/scripts/cleanup.py new file mode 100644 index 00000000..b768ad33 --- /dev/null +++ b/tools/ddb_migration/scripts/cleanup.py @@ -0,0 +1,153 @@ +"""Post-cutover cleanup of migration metadata on the target table. + +Run this after cutover is fully validated (typically 7-14 days post-cutover): + +* Sets a TTL attribute (``_ttl``) on every item with ``_tombstone=True`` so + DynamoDB's TTL feature deletes them. The target table must have TTL enabled + on the ``_ttl`` attribute. ``deploy.sh`` enables this automatically. +* Removes the ``_migration_ts`` attribute from regular items via paginated + scan-update. Idempotent: re-running is safe and does nothing on already-clean + items. + +Tombstone TTL defaults to 7 days from now. Override via ``--tombstone-ttl-days``. + +Exit codes +---------- + +* ``0`` — cleanup completed. +* ``1`` — at least one item failed to update. +* ``2`` — usage error. +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from typing import Any + +import boto3 +from botocore.exceptions import ClientError + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger("cleanup") + + +def iter_items(table: Any, projection: list[str], filter_expr: str | None = None) -> Any: + """Yield items from a paginated Scan using ProjectionExpression.""" + placeholders = {f"#a{i}": name for i, name in enumerate(projection)} + scan_kwargs: dict[str, Any] = { + "ProjectionExpression": ", ".join(placeholders.keys()), + "ExpressionAttributeNames": placeholders, + } + if filter_expr: + scan_kwargs["FilterExpression"] = filter_expr + while True: + resp = table.scan(**scan_kwargs) + for item in resp.get("Items", []): + yield item + if "LastEvaluatedKey" not in resp: + return + scan_kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +def expire_tombstones( + table: Any, + partition_key: str, + sort_key: str | None, + tombstone_ttl_seconds: int, +) -> dict[str, int]: + """Set _ttl on every tombstone item so DynamoDB TTL deletes it.""" + counts = {"updated": 0, "errors": 0} + expire_at = int(time.time()) + tombstone_ttl_seconds + proj = [partition_key] + if sort_key: + proj.append(sort_key) + proj.append("_tombstone") + for item in iter_items(table, proj, filter_expr="attribute_exists(#a2)"): + key = {partition_key: item[partition_key]} + if sort_key and sort_key in item: + key[sort_key] = item[sort_key] + try: + table.update_item( + Key=key, + UpdateExpression="SET #ttl = :ttl", + ExpressionAttributeNames={"#ttl": "_ttl"}, + ExpressionAttributeValues={":ttl": expire_at}, + ) + counts["updated"] += 1 + except ClientError as e: + log.exception("Failed to set TTL on %s: %s", key, e) + counts["errors"] += 1 + log.info("Tombstone TTL pass: updated=%d errors=%d", counts["updated"], counts["errors"]) + return counts + + +def remove_migration_ts(table: Any, partition_key: str, sort_key: str | None) -> dict[str, int]: + """Strip _migration_ts attribute from items that still carry it.""" + counts = {"updated": 0, "skipped": 0, "errors": 0} + proj = [partition_key] + if sort_key: + proj.append(sort_key) + proj.append("_migration_ts") + for item in iter_items(table, proj, filter_expr="attribute_exists(#a2)"): + key = {partition_key: item[partition_key]} + if sort_key and sort_key in item: + key[sort_key] = item[sort_key] + try: + table.update_item( + Key=key, + UpdateExpression="REMOVE #ts", + ExpressionAttributeNames={"#ts": "_migration_ts"}, + ConditionExpression="attribute_exists(#ts)", + ) + counts["updated"] += 1 + except ClientError as e: + if e.response["Error"]["Code"] == "ConditionalCheckFailedException": + counts["skipped"] += 1 + continue + log.exception("Failed to remove _migration_ts from %s: %s", key, e) + counts["errors"] += 1 + log.info( + "_migration_ts pass: updated=%d skipped=%d errors=%d", + counts["updated"], counts["skipped"], counts["errors"], + ) + return counts + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--target-table", default=os.environ.get("TARGET_TABLE")) + p.add_argument("--partition-key", default=os.environ.get("PARTITION_KEY", "pk")) + p.add_argument("--sort-key", default=os.environ.get("SORT_KEY")) + p.add_argument("--tombstone-ttl-days", type=int, default=7) + p.add_argument("--skip-tombstones", action="store_true") + p.add_argument("--skip-migration-ts", action="store_true") + return p.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if not args.target_table: + log.error("TARGET_TABLE is required") + return 2 + table = boto3.resource("dynamodb", region_name=args.region).Table(args.target_table) + total_errors = 0 + if not args.skip_tombstones: + total_errors += expire_tombstones( + table, args.partition_key, args.sort_key, args.tombstone_ttl_days * 86400, + )["errors"] + if not args.skip_migration_ts: + total_errors += remove_migration_ts(table, args.partition_key, args.sort_key)["errors"] + if total_errors > 0: + log.error("Cleanup completed with %d errors", total_errors) + return 1 + log.info("CLEANUP COMPLETE") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/scripts/convergence_check.py b/tools/ddb_migration/scripts/convergence_check.py new file mode 100644 index 00000000..65953fdb --- /dev/null +++ b/tools/ddb_migration/scripts/convergence_check.py @@ -0,0 +1,216 @@ +"""Pre-cutover convergence gate. + +Runs three checks in sequence and exits 0 only when all pass: + +1. **Iterator age** — stream-replay Lambda's ``IteratorAge`` is below + ``--max-iterator-age-ms`` (default 1000 ms). Polls until the timeout. +2. **DLQ depth** — the Lambda's failure DLQ is empty (visible + in-flight). +3. **Item count drift** — Scan COUNT (not the table-metadata ``ItemCount``, + which is updated only every ~6 hours) on both tables agrees within + ``--count-drift-pct`` (default 0.5%). ``--ignore-count-drift`` to skip. + +Exit codes +---------- + +* ``0`` — all checks passed; safe to cut over. +* ``1`` — at least one check failed. +* ``2`` — usage error (missing required env var or arg). +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from datetime import datetime, timedelta, timezone +from typing import Any + +import boto3 +from botocore.exceptions import ClientError + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger("convergence_check") + + +def check_iterator_age( + cw_client: Any, + lambda_function_name: str, + max_iterator_age_ms: int, + max_wait_seconds: int, + poll_interval: int = 10, + idle_grace_seconds: int = 120, +) -> bool: + """Poll IteratorAge until below threshold or until ``max_wait_seconds`` elapses. + + Lambda only emits IteratorAge on invocation. If the source table is idle, + the metric has no datapoints — that means caught up, not lagging. After + ``idle_grace_seconds`` of no datapoints we treat the absence as a pass. + """ + deadline = time.time() + max_wait_seconds + no_data_since: float | None = None + while time.time() < deadline: + end = datetime.now(timezone.utc) + start = end - timedelta(minutes=2) + try: + resp = cw_client.get_metric_statistics( + Namespace="AWS/Lambda", + MetricName="IteratorAge", + Dimensions=[{"Name": "FunctionName", "Value": lambda_function_name}], + StartTime=start, + EndTime=end, + Period=60, + Statistics=["Maximum"], + ) + except ClientError as e: + log.error("CloudWatch query failed: %s", e.response["Error"]["Code"]) + return False + points = resp.get("Datapoints", []) + if not points: + if no_data_since is None: + no_data_since = time.time() + log.info("No IteratorAge datapoints; starting %ds idle grace window", idle_grace_seconds) + elif time.time() - no_data_since >= idle_grace_seconds: + log.info("No IteratorAge datapoints for %ds — treating Lambda as idle/caught-up", idle_grace_seconds) + return True + time.sleep(poll_interval) + continue + no_data_since = None + max_age = max(p["Maximum"] for p in points) + log.info("IteratorAge max=%.0f ms (threshold %d ms)", max_age, max_iterator_age_ms) + if max_age <= max_iterator_age_ms: + return True + time.sleep(poll_interval) + log.error("IteratorAge did not converge within %ds", max_wait_seconds) + return False + + +def check_dlq_empty(sqs_client: Any, dlq_url: str) -> bool: + try: + resp = sqs_client.get_queue_attributes( + QueueUrl=dlq_url, + AttributeNames=["ApproximateNumberOfMessages", "ApproximateNumberOfMessagesNotVisible"], + ) + except ClientError as e: + log.error("SQS query failed: %s", e.response["Error"]["Code"]) + return False + visible = int(resp["Attributes"]["ApproximateNumberOfMessages"]) + in_flight = int(resp["Attributes"]["ApproximateNumberOfMessagesNotVisible"]) + log.info("DLQ depth: visible=%d in_flight=%d", visible, in_flight) + return visible == 0 and in_flight == 0 + + +def scan_count(ddb_client: Any, table_name: str, exclude_tombstones: bool = False) -> int: + """Authoritative item count via Scan with Select=COUNT (paginated). + + ``exclude_tombstones=True`` filters out items where ``_tombstone`` is set, + which is the right comparison for the target table (tombstones are + placeholders for deleted source items, not live data). + """ + total = 0 + kwargs: dict[str, Any] = {"TableName": table_name, "Select": "COUNT"} + if exclude_tombstones: + kwargs["FilterExpression"] = "attribute_not_exists(#t)" + kwargs["ExpressionAttributeNames"] = {"#t": "_tombstone"} + while True: + resp = ddb_client.scan(**kwargs) + total += resp.get("Count", 0) + if "LastEvaluatedKey" not in resp: + return total + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +def check_item_counts( + ddb_client: Any, + source_table: str, + target_table: str, + drift_pct: float, +) -> bool: + """Compare Scan COUNT on source vs (target − tombstones). + + Tombstones live in the target until ``cleanup.py`` runs post-cutover, so we + must exclude them when comparing logical item counts. + """ + log.info("Scanning source table %s (this may take several minutes)...", source_table) + src_count = scan_count(ddb_client, source_table) + log.info("Scanning target table %s (excluding tombstones)...", target_table) + tgt_count = scan_count(ddb_client, target_table, exclude_tombstones=True) + if src_count == 0: + log.warning("Source has 0 items; treating count check as vacuously passed") + return True + drift = abs(src_count - tgt_count) / src_count + log.info("Counts: source=%d target_live=%d drift=%.4f (max %.4f)", src_count, tgt_count, drift, drift_pct) + return drift <= drift_pct + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--source-table", default=os.environ.get("SOURCE_TABLE")) + p.add_argument("--target-table", default=os.environ.get("TARGET_TABLE")) + p.add_argument("--dlq-url", default=os.environ.get("DLQ_URL")) + p.add_argument( + "--lambda-function-name", + default=os.environ.get("LAMBDA_FUNCTION_NAME", "ddb-migration-stream-replay"), + ) + p.add_argument("--max-iterator-age-ms", type=int, default=1000) + p.add_argument("--max-wait-seconds", type=int, default=600) + p.add_argument("--count-drift-pct", type=float, default=0.005) + p.add_argument("--ignore-count-drift", action="store_true") + p.add_argument("--skip-iterator-age", action="store_true", help="(testing only)") + p.add_argument("--skip-dlq", action="store_true", help="(testing only)") + return p.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if not args.source_table or not args.target_table: + log.error("SOURCE_TABLE and TARGET_TABLE are required") + return 2 + if not args.dlq_url and not args.skip_dlq: + log.error("DLQ_URL is required (or pass --skip-dlq)") + return 2 + + cw = boto3.client("cloudwatch", region_name=args.region) + sqs = boto3.client("sqs", region_name=args.region) + ddb = boto3.client("dynamodb", region_name=args.region) + + results = [] + if args.skip_iterator_age: + log.warning("Skipping iterator-age check (testing only)") + results.append(("iterator_age", True)) + else: + results.append( + ("iterator_age", check_iterator_age( + cw, args.lambda_function_name, args.max_iterator_age_ms, args.max_wait_seconds, + )), + ) + if args.skip_dlq: + log.warning("Skipping DLQ check (testing only)") + results.append(("dlq_empty", True)) + else: + results.append(("dlq_empty", check_dlq_empty(sqs, args.dlq_url))) + if args.ignore_count_drift: + log.warning("Skipping count check (--ignore-count-drift)") + results.append(("count_match", True)) + else: + results.append(( + "count_match", + check_item_counts(ddb, args.source_table, args.target_table, args.count_drift_pct), + )) + + log.info("=" * 60) + for name, passed in results: + log.info(" %s %s", "PASS" if passed else "FAIL", name) + log.info("=" * 60) + + if all(passed for _, passed in results): + log.info("CONVERGENCE OK — proceed with cutover") + return 0 + log.error("CONVERGENCE FAILED — do not cut over") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/scripts/verify_cutover.py b/tools/ddb_migration/scripts/verify_cutover.py new file mode 100644 index 00000000..3890fcb6 --- /dev/null +++ b/tools/ddb_migration/scripts/verify_cutover.py @@ -0,0 +1,144 @@ +"""Sample-based source ↔ target cutover verifier. + +Picks N random items from the source table, looks up the corresponding key in +the target table, applies ``transform()`` to the source item, and asserts +deep-equality (ignoring migration metadata attributes). + +This is a confidence check, not a proof of completeness — for that, run +``convergence_check.py`` plus a full Scan COUNT comparison. Use this script +during a smoke test or as a CI gate before flipping app routing. + +Exit codes +---------- + +* ``0`` — sample matches. +* ``1`` — at least one mismatch or missing item. +* ``2`` — usage error. +""" + +from __future__ import annotations + +import argparse +import importlib +import logging +import os +import random +import sys +from typing import Any, Callable + +import boto3 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger("verify_cutover") + + +def load_transform() -> Callable[..., Any]: + module_name = os.environ.get("TRANSFORM_MODULE") + if module_name: + try: + mod = importlib.import_module(module_name) + return getattr(mod, "transform") + except (ImportError, AttributeError) as e: + log.warning("transform module %s failed to load: %s", module_name, e) + try: + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + from transform import transform # type: ignore + + return transform + except ImportError: + return lambda item, source_event=None: item + + +def sample_source_items(source_table: Any, sample_size: int) -> list[dict[str, Any]]: + """Reservoir-sample items from the source table via Scan.""" + reservoir: list[dict[str, Any]] = [] + seen = 0 + scan_kwargs: dict[str, Any] = {} + while True: + resp = source_table.scan(**scan_kwargs) + for item in resp.get("Items", []): + seen += 1 + if len(reservoir) < sample_size: + reservoir.append(item) + else: + idx = random.randint(0, seen - 1) + if idx < sample_size: + reservoir[idx] = item + if "LastEvaluatedKey" not in resp: + break + scan_kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + log.info("Sampled %d of %d source items", len(reservoir), seen) + return reservoir + + +def items_match(source_item: dict[str, Any], target_item: dict[str, Any]) -> bool: + """Deep-equality ignoring migration metadata.""" + s = {k: v for k, v in source_item.items() if not k.startswith("_migration") and not k.startswith("_tombstone") and not k.startswith("_ttl")} + t = {k: v for k, v in target_item.items() if not k.startswith("_migration") and not k.startswith("_tombstone") and not k.startswith("_ttl")} + return s == t + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument("--region", default=os.environ.get("REGION", "us-east-1")) + p.add_argument("--source-table", default=os.environ.get("SOURCE_TABLE")) + p.add_argument("--target-table", default=os.environ.get("TARGET_TABLE")) + p.add_argument("--partition-key", default=os.environ.get("PARTITION_KEY", "pk")) + p.add_argument("--sort-key", default=os.environ.get("SORT_KEY")) + p.add_argument("--sample-size", type=int, default=1000) + return p.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if not args.source_table or not args.target_table: + log.error("SOURCE_TABLE and TARGET_TABLE are required") + return 2 + transform = load_transform() + ddb = boto3.resource("dynamodb", region_name=args.region) + source = ddb.Table(args.source_table) + target = ddb.Table(args.target_table) + + samples = sample_source_items(source, args.sample_size) + if not samples: + log.warning("Source table is empty; nothing to verify") + return 0 + + matched = 0 + missing = 0 + diverged = 0 + for src_item in samples: + key: dict[str, Any] = {args.partition_key: src_item[args.partition_key]} + if args.sort_key and args.sort_key in src_item: + key[args.sort_key] = src_item[args.sort_key] + resp = target.get_item(Key=key) + target_item = resp.get("Item") + if not target_item: + missing += 1 + log.warning("Missing in target: %s", key) + continue + expected = transform(dict(src_item)) + if expected is None: + # Source item filtered by transform; target should also lack it. + if target_item: + diverged += 1 + log.warning("Filtered source but target has item: %s", key) + continue + if items_match(expected, target_item): + matched += 1 + else: + diverged += 1 + log.warning("Diverged: %s", key) + + log.info("=" * 60) + log.info(" matched=%d missing=%d diverged=%d total=%d", matched, missing, diverged, len(samples)) + log.info("=" * 60) + if missing == 0 and diverged == 0: + log.info("VERIFY OK") + return 0 + log.error("VERIFY FAILED") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/ddb_migration/teardown.sh b/tools/ddb_migration/teardown.sh new file mode 100755 index 00000000..95711a52 --- /dev/null +++ b/tools/ddb_migration/teardown.sh @@ -0,0 +1,88 @@ +#!/usr/bin/env bash +# Removes everything deploy.sh provisioned. Best-effort: prints anything it +# could not remove. Refuses to run unless CONFIRM=yes. +# +# Does NOT delete: +# - the source table +# - the target table (you must verify cutover and clean it up yourself) +# - data in the S3 export bucket (deletes the bucket only after empty) + +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [[ -f "$SCRIPT_DIR/config.env" ]]; then + # shellcheck source=/dev/null + source "$SCRIPT_DIR/config.env" +fi + +if [[ "${CONFIRM:-}" != "yes" ]]; then + cat >&2 </dev/null || { printf " (skipped/failed: %s)\n" "$*" >&2; errs=$((errs+1)); }; } + +step "Deleting event source mappings for $LAMBDA_FUNCTION_NAME" +mapfile -t esms < <(aws lambda list-event-source-mappings \ + --function-name "$LAMBDA_FUNCTION_NAME" \ + --region "$REGION" \ + --query 'EventSourceMappings[].UUID' --output text 2>/dev/null | tr '\t' '\n') +for esm in "${esms[@]}"; do + [[ -z "$esm" ]] && continue + try aws lambda delete-event-source-mapping --uuid "$esm" --region "$REGION" +done + +step "Deleting Lambda $LAMBDA_FUNCTION_NAME" +try aws lambda delete-function --function-name "$LAMBDA_FUNCTION_NAME" --region "$REGION" + +step "Deleting IAM role $LAMBDA_ROLE_NAME" +try aws iam delete-role-policy --role-name "$LAMBDA_ROLE_NAME" --policy-name "ddb-migration-stream-replay-inline" +try aws iam delete-role --role-name "$LAMBDA_ROLE_NAME" + +step "Deleting CloudWatch alarms" +try aws cloudwatch delete-alarms \ + --alarm-names ddb-migration-iterator-age-warning ddb-migration-iterator-age-critical \ + --region "$REGION" + +step "Deleting SNS topic $SNS_TOPIC_NAME" +sns_arn="$(aws sns list-topics --region "$REGION" --query "Topics[?ends_with(TopicArn, ':$SNS_TOPIC_NAME')].TopicArn | [0]" --output text 2>/dev/null)" +if [[ "$sns_arn" != "None" && -n "$sns_arn" ]]; then + try aws sns delete-topic --topic-arn "$sns_arn" --region "$REGION" +fi + +step "Deleting SQS DLQ $DLQ_NAME" +dlq_url="$(aws sqs get-queue-url --queue-name "$DLQ_NAME" --region "$REGION" --query 'QueueUrl' --output text 2>/dev/null || echo "")" +if [[ -n "$dlq_url" && "$dlq_url" != "None" ]]; then + try aws sqs delete-queue --queue-url "$dlq_url" --region "$REGION" +fi + +step "Deleting export bucket s3://$EXPORT_BUCKET (must be empty)" +try aws s3api delete-bucket --bucket "$EXPORT_BUCKET" --region "$REGION" + +if [[ $errs -gt 0 ]]; then + printf "\n[teardown] completed with %d step(s) skipped or failed; see warnings above\n" "$errs" + exit 1 +fi +printf "\n[teardown] complete\n" diff --git a/tools/ddb_migration/tests/__init__.py b/tools/ddb_migration/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/ddb_migration/tests/conftest.py b/tools/ddb_migration/tests/conftest.py new file mode 100644 index 00000000..d8b37017 --- /dev/null +++ b/tools/ddb_migration/tests/conftest.py @@ -0,0 +1,67 @@ +"""Shared test fixtures. + +We use moto for DynamoDB / S3 / SQS / CloudWatch mocking. Each fixture is +function-scoped so tests get an isolated AWS environment. +""" + +from __future__ import annotations + +import os + +import boto3 +import pytest +from moto import mock_aws + + +@pytest.fixture(autouse=True) +def aws_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + """Force boto3 to use dummy credentials so it doesn't read ~/.aws/.""" + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") + monkeypatch.delenv("AWS_PROFILE", raising=False) + + +@pytest.fixture +def aws(): + """Activate moto for the duration of one test.""" + with mock_aws(): + yield + + +@pytest.fixture +def target_table(aws): + """Create a hash+range target table mirroring the demo schema.""" + ddb = boto3.resource("dynamodb", region_name="us-east-1") + ddb.create_table( + TableName="target", + AttributeDefinitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + ], + KeySchema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + return ddb.Table("target") + + +@pytest.fixture +def source_table(aws): + ddb = boto3.resource("dynamodb", region_name="us-east-1") + ddb.create_table( + TableName="source", + AttributeDefinitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + ], + KeySchema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + BillingMode="PAY_PER_REQUEST", + ) + return ddb.Table("source") diff --git a/tools/ddb_migration/tests/test_backfill.py b/tools/ddb_migration/tests/test_backfill.py new file mode 100644 index 00000000..6e8d4b18 --- /dev/null +++ b/tools/ddb_migration/tests/test_backfill.py @@ -0,0 +1,126 @@ +"""Tests for scripts/backfill.py.""" + +from __future__ import annotations + +import gzip +import io +import json +from typing import Any +from unittest.mock import MagicMock + +import boto3 +import pytest + +import backfill + + +def upload_export(bucket: str, prefix: str, items: list[dict]) -> None: + """Lay down a fake DynamoDB S3 export: manifest-summary, manifest-files, data file.""" + s3 = boto3.client("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket) + export_dir = f"{prefix.rstrip('/')}/AWSDynamoDB/01234567890-deadbeef/" + summary = {"manifestFilesS3Key": f"{export_dir}manifest-files.json"} + s3.put_object(Bucket=bucket, Key=f"{export_dir}manifest-summary.json", Body=json.dumps(summary)) + data_key = f"{export_dir}data/file1.json.gz" + manifest_lines = [json.dumps({"dataFileS3Key": data_key})] + s3.put_object(Bucket=bucket, Key=f"{export_dir}manifest-files.json", Body="\n".join(manifest_lines)) + raw = "\n".join(json.dumps({"Item": item}) for item in items).encode() + buf = io.BytesIO() + with gzip.GzipFile(fileobj=buf, mode="w") as gz: + gz.write(raw) + s3.put_object(Bucket=bucket, Key=data_key, Body=buf.getvalue()) + + +def ddb_serialize(item: dict) -> dict: + """Tiny helper: convert a plain Python dict to DDB-JSON for the export.""" + out: dict = {} + for k, v in item.items(): + if isinstance(v, str): + out[k] = {"S": v} + elif isinstance(v, bool): + out[k] = {"BOOL": v} + elif isinstance(v, (int, float)): + out[k] = {"N": str(v)} + return out + + +def test_get_data_file_keys_parses_manifest(aws: Any) -> None: + items = [{"pk": "a", "sk": "1"}, {"pk": "b", "sk": "1"}] + upload_export("bkt", "exports/", [ddb_serialize(i) for i in items]) + s3 = boto3.client("s3", region_name="us-east-1") + keys = backfill.get_data_file_keys(s3, "bkt", "exports/") + assert len(keys) == 1 + assert keys[0].endswith("data/file1.json.gz") + + +def test_put_item_with_retry_writes_with_zero_migration_ts(target_table: Any) -> None: + outcome = backfill.put_item_with_retry(target_table, {"pk": "a", "sk": "1", "x": 1}, "pk") + assert outcome == "written" + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert int(item["_migration_ts"]) == 0 + + +def test_put_item_with_retry_skips_when_newer_exists(target_table: Any) -> None: + target_table.put_item(Item={"pk": "a", "sk": "1", "_migration_ts": 500}) + outcome = backfill.put_item_with_retry(target_table, {"pk": "a", "sk": "1"}, "pk") + assert outcome == "skipped" + + +def test_should_pause_true_when_iterator_age_above_threshold() -> None: + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": [{"Maximum": 70_000_000}]} + assert backfill.should_pause(cw, "fn", pause_threshold_hours=18) is True + + +def test_should_pause_false_when_iterator_age_below_threshold() -> None: + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": [{"Maximum": 1000}]} + assert backfill.should_pause(cw, "fn", pause_threshold_hours=18) is False + + +def test_should_pause_false_when_no_datapoints() -> None: + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": []} + assert backfill.should_pause(cw, "fn", pause_threshold_hours=18) is False + + +def test_dry_run_does_not_write(aws: Any, target_table: Any) -> None: + items = [{"pk": "a", "sk": "1"}, {"pk": "b", "sk": "1"}] + upload_export("bkt", "exports/", [ddb_serialize(i) for i in items]) + s3 = boto3.client("s3", region_name="us-east-1") + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": []} + config = { + "transform": lambda item, source_event=None: item, + "partition_key": "pk", + "lambda_function_name": "fn", + "pause_threshold_hours": 18, + "circuit_breaker_check_interval": 5000, + "dry_run": True, + } + keys = backfill.get_data_file_keys(s3, "bkt", "exports/") + counts = backfill.process_data_file(s3, cw, target_table, "bkt", keys[0], config) + assert counts["items"] == 2 + assert counts["written"] == 0 + assert target_table.scan()["Count"] == 0 + + +def test_full_file_pipeline_writes_items(aws: Any, target_table: Any) -> None: + items = [{"pk": f"a{i}", "sk": "1", "v": i} for i in range(3)] + upload_export("bkt", "exports/", [ddb_serialize(i) for i in items]) + s3 = boto3.client("s3", region_name="us-east-1") + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": []} + config = { + "transform": lambda item, source_event=None: item, + "partition_key": "pk", + "lambda_function_name": "fn", + "pause_threshold_hours": 18, + "circuit_breaker_check_interval": 5000, + "dry_run": False, + } + keys = backfill.get_data_file_keys(s3, "bkt", "exports/") + counts = backfill.process_data_file(s3, cw, target_table, "bkt", keys[0], config) + assert counts["items"] == 3 + assert counts["written"] == 3 + assert target_table.scan()["Count"] == 3 diff --git a/tools/ddb_migration/tests/test_cleanup.py b/tools/ddb_migration/tests/test_cleanup.py new file mode 100644 index 00000000..a2c5596a --- /dev/null +++ b/tools/ddb_migration/tests/test_cleanup.py @@ -0,0 +1,48 @@ +"""Tests for scripts/cleanup.py.""" + +from __future__ import annotations + +import time +from typing import Any + +import cleanup + + +def test_expire_tombstones_sets_ttl(target_table: Any) -> None: + target_table.put_item(Item={"pk": "a", "sk": "1", "_tombstone": True, "_migration_ts": 100}) + target_table.put_item(Item={"pk": "b", "sk": "1", "_tombstone": True, "_migration_ts": 200}) + target_table.put_item(Item={"pk": "c", "sk": "1", "name": "live"}) + + counts = cleanup.expire_tombstones(target_table, "pk", "sk", tombstone_ttl_seconds=86400) + assert counts["updated"] == 2 + assert counts["errors"] == 0 + + a = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert "_ttl" in a + assert a["_ttl"] >= int(time.time()) + + c = target_table.get_item(Key={"pk": "c", "sk": "1"})["Item"] + assert "_ttl" not in c + + +def test_remove_migration_ts_strips_attribute(target_table: Any) -> None: + target_table.put_item(Item={"pk": "a", "sk": "1", "_migration_ts": 100, "name": "x"}) + target_table.put_item(Item={"pk": "b", "sk": "1", "name": "y"}) + + counts = cleanup.remove_migration_ts(target_table, "pk", "sk") + assert counts["updated"] == 1 + + a = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert "_migration_ts" not in a + assert a["name"] == "x" + + b = target_table.get_item(Key={"pk": "b", "sk": "1"})["Item"] + assert "_migration_ts" not in b + + +def test_idempotent_rerun(target_table: Any) -> None: + target_table.put_item(Item={"pk": "a", "sk": "1", "_migration_ts": 100}) + cleanup.remove_migration_ts(target_table, "pk", "sk") + counts = cleanup.remove_migration_ts(target_table, "pk", "sk") + # Already removed; nothing to update. + assert counts["updated"] == 0 diff --git a/tools/ddb_migration/tests/test_convergence_check.py b/tools/ddb_migration/tests/test_convergence_check.py new file mode 100644 index 00000000..5a74cbcc --- /dev/null +++ b/tools/ddb_migration/tests/test_convergence_check.py @@ -0,0 +1,126 @@ +"""Tests for scripts/convergence_check.py.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import boto3 +import pytest + +import convergence_check + + +def test_iterator_age_passes_when_below_threshold() -> None: + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": [{"Maximum": 100}]} + assert convergence_check.check_iterator_age(cw, "fn", max_iterator_age_ms=1000, max_wait_seconds=2, poll_interval=0) is True + + +def test_iterator_age_fails_when_above_threshold() -> None: + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": [{"Maximum": 100_000}]} + assert convergence_check.check_iterator_age(cw, "fn", max_iterator_age_ms=1000, max_wait_seconds=1, poll_interval=0) is False + + +def test_dlq_empty_passes_at_zero_depth() -> None: + sqs = MagicMock() + sqs.get_queue_attributes.return_value = { + "Attributes": {"ApproximateNumberOfMessages": "0", "ApproximateNumberOfMessagesNotVisible": "0"}, + } + assert convergence_check.check_dlq_empty(sqs, "https://example/dlq") is True + + +def test_dlq_empty_fails_when_visible() -> None: + sqs = MagicMock() + sqs.get_queue_attributes.return_value = { + "Attributes": {"ApproximateNumberOfMessages": "1", "ApproximateNumberOfMessagesNotVisible": "0"}, + } + assert convergence_check.check_dlq_empty(sqs, "https://example/dlq") is False + + +def test_dlq_empty_fails_when_in_flight() -> None: + sqs = MagicMock() + sqs.get_queue_attributes.return_value = { + "Attributes": {"ApproximateNumberOfMessages": "0", "ApproximateNumberOfMessagesNotVisible": "1"}, + } + assert convergence_check.check_dlq_empty(sqs, "https://example/dlq") is False + + +def test_count_check_uses_scan_not_metadata(aws: Any, source_table: Any, target_table: Any) -> None: + """Critical: gap #12 was that the original PR used describe-table ItemCount (~6h stale). + This test fails if a regression replaces the Scan with describe_table. + """ + for i in range(5): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1"}) + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "_migration_ts": 0}) + ddb = boto3.client("dynamodb", region_name="us-east-1") + assert convergence_check.check_item_counts(ddb, "source", "target", drift_pct=0.005) is True + + +def test_count_check_fails_above_drift_threshold(aws: Any, source_table: Any, target_table: Any) -> None: + for i in range(100): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1"}) + for i in range(50): + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "_migration_ts": 0}) + ddb = boto3.client("dynamodb", region_name="us-east-1") + # 50% drift, threshold 0.5%, must fail + assert convergence_check.check_item_counts(ddb, "source", "target", drift_pct=0.005) is False + + +def test_count_check_passes_when_drift_within_tolerance(aws: Any, source_table: Any, target_table: Any) -> None: + for i in range(1000): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1"}) + for i in range(998): + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "_migration_ts": 0}) + ddb = boto3.client("dynamodb", region_name="us-east-1") + # 0.2% drift, threshold 0.5%, must pass + assert convergence_check.check_item_counts(ddb, "source", "target", drift_pct=0.005) is True + + +def test_main_skip_flags_keep_zero_exit(aws: Any, source_table: Any, target_table: Any) -> None: + source_table.put_item(Item={"pk": "a", "sk": "1"}) + target_table.put_item(Item={"pk": "a", "sk": "1", "_migration_ts": 0}) + rc = convergence_check.main([ + "--source-table", "source", + "--target-table", "target", + "--skip-iterator-age", + "--skip-dlq", + ]) + assert rc == 0 + + +def test_main_returns_2_when_required_args_missing() -> None: + rc = convergence_check.main([]) + assert rc == 2 + + +def test_count_check_excludes_tombstones_from_target( + aws: Any, source_table: Any, target_table: Any +) -> None: + """Tombstones are placeholders for deleted source items — excluded from the live count. + + Without this filter, deleting an item from source then having stream-replay write a + tombstone to target produces source_count == N-1 vs target_count == N, a false drift. + """ + for i in range(100): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1"}) + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "_migration_ts": 0}) + # Simulate 5 deletes that became tombstones in target. + for i in range(95, 100): + source_table.delete_item(Key={"pk": f"k{i}", "sk": "1"}) + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "_tombstone": True, "_migration_ts": 100}) + ddb = boto3.client("dynamodb", region_name="us-east-1") + # Without tombstone exclusion: source=95, target=100 → 5.3% drift, would fail. + # With tombstone exclusion: source=95, target_live=95 → 0% drift, passes. + assert convergence_check.check_item_counts(ddb, "source", "target", drift_pct=0.005) is True + + +def test_iterator_age_treats_no_data_as_pass_after_grace(monkeypatch: Any) -> None: + """Idle Lambda emits no IteratorAge metric — must not block forever.""" + cw = MagicMock() + cw.get_metric_statistics.return_value = {"Datapoints": []} + # idle_grace_seconds=0 means first no-data hit immediately passes. + assert convergence_check.check_iterator_age( + cw, "fn", max_iterator_age_ms=1000, max_wait_seconds=2, poll_interval=0, idle_grace_seconds=0, + ) is True diff --git a/tools/ddb_migration/tests/test_stream_replay.py b/tools/ddb_migration/tests/test_stream_replay.py new file mode 100644 index 00000000..0dc74858 --- /dev/null +++ b/tools/ddb_migration/tests/test_stream_replay.py @@ -0,0 +1,121 @@ +"""Tests for lambda/stream_replay.py. + +moto's stream-record handling is partial, so we drive the handler with +hand-rolled events and let it write to a moto-backed target table. +""" + +from __future__ import annotations + +import importlib +import sys +from typing import Any + +import boto3 +import pytest + + +@pytest.fixture +def stream_replay_module(monkeypatch: pytest.MonkeyPatch, target_table: Any): + """Import lambda/stream_replay.py with the right env vars set.""" + monkeypatch.setenv("TARGET_TABLE", "target") + monkeypatch.setenv("PARTITION_KEY", "pk") + monkeypatch.setenv("TARGET_REGION", "us-east-1") + sys.modules.pop("stream_replay", None) + mod = importlib.import_module("stream_replay") + return mod + + +def make_insert_event(pk: str, sk: str, ts: float, **fields: Any) -> dict: + new_image = {"pk": {"S": pk}, "sk": {"S": sk}, **{k: {"S": str(v)} for k, v in fields.items()}} + return { + "Records": [{ + "eventID": f"e-{pk}-{sk}", + "eventName": "INSERT", + "dynamodb": { + "ApproximateCreationDateTime": ts, + "Keys": {"pk": {"S": pk}, "sk": {"S": sk}}, + "NewImage": new_image, + }, + }] + } + + +def make_remove_event(pk: str, sk: str, ts: float) -> dict: + return { + "Records": [{ + "eventID": f"e-{pk}-{sk}", + "eventName": "REMOVE", + "dynamodb": { + "ApproximateCreationDateTime": ts, + "Keys": {"pk": {"S": pk}, "sk": {"S": sk}}, + "OldImage": {"pk": {"S": pk}, "sk": {"S": sk}, "status": {"S": "PAID"}}, + }, + }] + } + + +def test_insert_writes_item_with_migration_ts(stream_replay_module, target_table) -> None: + result = stream_replay_module.handler(make_insert_event("a", "1", 100.0, status="NEW"), None) + assert result == {"batchItemFailures": []} + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert item["status"] == "NEW" + assert float(item["_migration_ts"]) == 100.0 + + +def test_newer_event_overwrites_older(stream_replay_module, target_table) -> None: + stream_replay_module.handler(make_insert_event("a", "1", 100.0, status="NEW"), None) + stream_replay_module.handler(make_insert_event("a", "1", 200.0, status="PAID"), None) + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert item["status"] == "PAID" + + +def test_older_event_does_not_overwrite_newer(stream_replay_module, target_table) -> None: + stream_replay_module.handler(make_insert_event("a", "1", 200.0, status="PAID"), None) + stream_replay_module.handler(make_insert_event("a", "1", 100.0, status="NEW"), None) + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert item["status"] == "PAID" + + +def test_remove_writes_tombstone(stream_replay_module, target_table) -> None: + stream_replay_module.handler(make_remove_event("a", "1", 300.0), None) + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert item["_tombstone"] is True + assert "status" not in item # full image not persisted; only key + flag + + +def test_unhandled_exception_reported_as_batch_failure( + monkeypatch: pytest.MonkeyPatch, stream_replay_module +) -> None: + """Errors other than ConditionalCheckFailed get appended to batchItemFailures.""" + + def boom(*a, **kw): + raise RuntimeError("boom") + + monkeypatch.setattr(stream_replay_module, "_conditional_put", boom) + event = make_insert_event("a", "1", 100.0) + result = stream_replay_module.handler(event, None) + assert result == {"batchItemFailures": [{"itemIdentifier": "e-a-1"}]} + + +def test_transform_module_env_var_loads_custom_module( + monkeypatch: pytest.MonkeyPatch, target_table, tmp_path +) -> None: + """TRANSFORM_MODULE must actually be honored — fixes gap #1 from the review.""" + custom = tmp_path / "my_transform.py" + custom.write_text( + "def transform(item, source_event=None):\n" + " item['custom'] = True\n" + " return item\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + monkeypatch.setenv("TARGET_TABLE", "target") + monkeypatch.setenv("PARTITION_KEY", "pk") + monkeypatch.setenv("TARGET_REGION", "us-east-1") + monkeypatch.setenv("TRANSFORM_MODULE", "my_transform") + + sys.modules.pop("stream_replay", None) + mod = importlib.import_module("stream_replay") + + mod.handler(make_insert_event("a", "1", 100.0, status="NEW"), None) + item = target_table.get_item(Key={"pk": "a", "sk": "1"})["Item"] + assert item["custom"] is True diff --git a/tools/ddb_migration/tests/test_transform.py b/tools/ddb_migration/tests/test_transform.py new file mode 100644 index 00000000..026b156e --- /dev/null +++ b/tools/ddb_migration/tests/test_transform.py @@ -0,0 +1,27 @@ +from transform import transform + + +def test_identity_passes_item_unchanged() -> None: + item = {"pk": "a", "sk": "1", "amount": 5} + assert transform(item) == item + + +def test_returns_none_means_skip() -> None: + """Caller contract: returning None signals 'skip this item.' + + The default identity transform never returns None, but a user-supplied + transform can. This test documents the contract via a sample. + """ + + def filter_deleted(item, source_event=None): + if item.get("status") == "DELETED": + return None + return item + + assert filter_deleted({"pk": "a", "status": "ACTIVE"}) is not None + assert filter_deleted({"pk": "a", "status": "DELETED"}) is None + + +def test_source_event_argument_is_optional() -> None: + transform({"pk": "a"}, source_event=None) + transform({"pk": "a"}, source_event={"eventName": "INSERT"}) diff --git a/tools/ddb_migration/tests/test_verify_cutover.py b/tools/ddb_migration/tests/test_verify_cutover.py new file mode 100644 index 00000000..14342e03 --- /dev/null +++ b/tools/ddb_migration/tests/test_verify_cutover.py @@ -0,0 +1,68 @@ +"""Tests for scripts/verify_cutover.py.""" + +from __future__ import annotations + +from typing import Any + +import verify_cutover + + +def test_items_match_ignores_migration_metadata() -> None: + src = {"pk": "a", "sk": "1", "name": "x"} + tgt = {"pk": "a", "sk": "1", "name": "x", "_migration_ts": 100} + assert verify_cutover.items_match(src, tgt) is True + + +def test_items_match_detects_data_divergence() -> None: + src = {"pk": "a", "sk": "1", "name": "x"} + tgt = {"pk": "a", "sk": "1", "name": "y"} + assert verify_cutover.items_match(src, tgt) is False + + +def test_main_succeeds_when_target_matches( + monkeypatch: Any, source_table: Any, target_table: Any +) -> None: + for i in range(20): + item = {"pk": f"k{i}", "sk": "1", "v": str(i)} + source_table.put_item(Item=item) + target_table.put_item(Item={**item, "_migration_ts": 100}) + rc = verify_cutover.main([ + "--source-table", "source", + "--target-table", "target", + "--partition-key", "pk", + "--sort-key", "sk", + "--sample-size", "20", + ]) + assert rc == 0 + + +def test_main_fails_when_target_missing_items( + source_table: Any, target_table: Any +) -> None: + for i in range(20): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1", "v": str(i)}) + # Only 5 in target. + for i in range(5): + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "v": str(i), "_migration_ts": 100}) + rc = verify_cutover.main([ + "--source-table", "source", + "--target-table", "target", + "--partition-key", "pk", + "--sort-key", "sk", + "--sample-size", "20", + ]) + assert rc == 1 + + +def test_main_fails_when_data_diverges(source_table: Any, target_table: Any) -> None: + for i in range(20): + source_table.put_item(Item={"pk": f"k{i}", "sk": "1", "v": str(i)}) + target_table.put_item(Item={"pk": f"k{i}", "sk": "1", "v": "WRONG", "_migration_ts": 100}) + rc = verify_cutover.main([ + "--source-table", "source", + "--target-table", "target", + "--partition-key", "pk", + "--sort-key", "sk", + "--sample-size", "20", + ]) + assert rc == 1 diff --git a/tools/ddb_migration/transform.py b/tools/ddb_migration/transform.py new file mode 100644 index 00000000..f6c2062e --- /dev/null +++ b/tools/ddb_migration/transform.py @@ -0,0 +1,65 @@ +"""Shared item transformation for DynamoDB zero-downtime migration. + +Both ``lambda/stream_replay.py`` and ``scripts/backfill.py`` import ``transform`` +from this module. Keeping a single implementation avoids divergence between the +backfill and live-replay paths — divergence breaks the conflict-resolution +invariant (newer ``_migration_ts`` wins) because the same logical item could be +written under two different shapes. + +To customize for a real migration, either: + +* edit ``transform`` below, or +* set the environment variable ``TRANSFORM_MODULE`` to a Python module path + (e.g. ``my_transforms.orders``) that exposes a ``transform(item, source_event=None)`` + function. The Lambda and backfill scripts will load it via ``importlib`` at + startup. The module must be importable from ``PYTHONPATH``. + +Contract: + +* Input ``item`` is a regular Python ``dict`` (DynamoDB JSON already deserialized). +* ``source_event`` is the full DynamoDB Streams event record for stream replay, + or ``None`` when called from the backfill path. +* Return the (possibly mutated) ``dict`` to write to the target. +* Return ``None`` to skip the item entirely (filter pattern). + +Examples +-------- + +Rename an attribute:: + + def transform(item, source_event=None): + if "old_name" in item: + item["new_name"] = item.pop("old_name") + return item + +Compute a new GSI key:: + + def transform(item, source_event=None): + item["status_date_idx"] = f"{item['status']}#{item['created_at']}" + return item + +Filter out items by status:: + + def transform(item, source_event=None): + if item.get("status") == "DELETED": + return None + return item + +Convert legacy ``Decimal`` floats:: + + from decimal import Decimal + + def transform(item, source_event=None): + if "price" in item and isinstance(item["price"], float): + item["price"] = Decimal(str(item["price"])) + return item +""" + +from __future__ import annotations + +from typing import Any + + +def transform(item: dict[str, Any], source_event: dict[str, Any] | None = None) -> dict[str, Any] | None: + """Default identity transform — pass items through unchanged.""" + return item