From e8a7da791391f6c65edb2bb2588c3cabd298780c Mon Sep 17 00:00:00 2001 From: col Date: Fri, 26 Jun 2026 11:05:04 -0700 Subject: [PATCH 01/22] ci: GitHub Actions CI pipeline for bulk_executor (#195) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a full CI pipeline that runs on every push to main and on PRs: ## Workflow chain (sequential via workflow_run triggers) unit tests (push+PR) → connector smoke → e2e fill/copy/update/delete/diff - Unit tests: Python 3.11, --ignore=tests/e2e, runs on PRs for fast feedback - Connector smoke: read-path validation (count/find/sql/load) via OIDC auth - E2e commands: one workflow per command, fan out in parallel after smoke - Flake retry: auto-reruns a failed e2e once; alerts only on 2nd failure ## Infrastructure scripts - e2e-bootstrap.sh: create OIDC provider + IAM role + secrets in one command - e2e-teardown.sh: clean removal of all CI infrastructure from an account - e2e-switch-account.sh: interactive teardown-old + bootstrap-new ## IAM permissions (least-privilege for the github-actions-e2e-runner role) - DynamoDB: CRUD on named test tables + Create/Delete for bulk-e2e-* transient tables - Glue: Start/Get job runs on bulk_dynamodb - CloudWatch Logs: DescribeLogGroups + StartLiveTail (CLI streams Glue output) - S3: read/write on aws-glue-bulk-dynamodb-* bucket - IAM: PassRole to glue.amazonaws.com for AWSGlueServiceRole* - STS: GetCallerIdentity (account guard) ## Security - All AWS credentials via OIDC (no long-lived keys) - Account ID and table names stored as GitHub Secrets - OIDC trust scoped to specific repos + main branch only - Security suite excluded from CI (requires admin, mutates shared job) Tested end-to-end on relentlesscol/amazon-dynamodb-tools fork — full pipeline green including all 5 e2e command tests. --- .github/scripts/e2e-bootstrap.sh | 238 ++++++++++++++++++ .github/scripts/e2e-switch-account.sh | 74 ++++++ .github/scripts/e2e-teardown.sh | 69 +++++ .../bulk-executor-connector-smoke.yml | 58 +++++ .github/workflows/bulk-executor-e2e-copy.yml | 58 +++++ .../workflows/bulk-executor-e2e-delete.yml | 58 +++++ .github/workflows/bulk-executor-e2e-diff.yml | 58 +++++ .github/workflows/bulk-executor-e2e-fill.yml | 58 +++++ .../workflows/bulk-executor-e2e-update.yml | 58 +++++ .../workflows/bulk-executor-flake-retry.yml | 64 +++++ .../workflows/bulk-executor-unit-tests.yml | 43 ++++ 11 files changed, 836 insertions(+) create mode 100755 .github/scripts/e2e-bootstrap.sh create mode 100755 .github/scripts/e2e-switch-account.sh create mode 100755 .github/scripts/e2e-teardown.sh create mode 100644 .github/workflows/bulk-executor-connector-smoke.yml create mode 100644 .github/workflows/bulk-executor-e2e-copy.yml create mode 100644 .github/workflows/bulk-executor-e2e-delete.yml create mode 100644 .github/workflows/bulk-executor-e2e-diff.yml create mode 100644 .github/workflows/bulk-executor-e2e-fill.yml create mode 100644 .github/workflows/bulk-executor-e2e-update.yml create mode 100644 .github/workflows/bulk-executor-flake-retry.yml create mode 100644 .github/workflows/bulk-executor-unit-tests.yml diff --git a/.github/scripts/e2e-bootstrap.sh b/.github/scripts/e2e-bootstrap.sh new file mode 100755 index 00000000..cbd10b4c --- /dev/null +++ b/.github/scripts/e2e-bootstrap.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Bootstrap GitHub Actions e2e infrastructure in a given AWS account. +# Creates: OIDC identity provider, IAM role, inline policy, GitHub secrets. +# +# Prerequisites: +# - AWS CLI authenticated to the target account (Admin role) +# - gh CLI authenticated with repo admin access +# - jq installed +# +# Usage: +# .github/scripts/e2e-bootstrap.sh \ +# --account-id 123456789012 \ +# --region us-east-1 \ +# --read-table tiny-boat \ +# --write-table mini-boat \ +# --repos "awslabs/amazon-dynamodb-tools,relentlesscol/amazon-dynamodb-tools" + +ROLE_NAME="github-actions-e2e-runner" +POLICY_NAME="e2e-test-access" + +usage() { + echo "Usage: $0 --account-id ID --region REGION --read-table TABLE --write-table TABLE --repos REPO1,REPO2" + exit 1 +} + +while [[ $# -gt 0 ]]; do + case $1 in + --account-id) ACCOUNT_ID="$2"; shift 2 ;; + --region) REGION="$2"; shift 2 ;; + --read-table) READ_TABLE="$2"; shift 2 ;; + --write-table) WRITE_TABLE="$2"; shift 2 ;; + --repos) IFS=',' read -ra REPOS <<< "$2"; shift 2 ;; + *) usage ;; + esac +done + +[[ -z "${ACCOUNT_ID:-}" || -z "${REGION:-}" || -z "${READ_TABLE:-}" || -z "${WRITE_TABLE:-}" || ${#REPOS[@]} -eq 0 ]] && usage + +echo "==> Bootstrapping e2e CI in account ${ACCOUNT_ID} (${REGION})" + +# --- OIDC Provider --- +OIDC_ARN="arn:aws:iam::${ACCOUNT_ID}:oidc-provider/token.actions.githubusercontent.com" +if aws iam get-open-id-connect-provider --open-id-connect-provider-arn "${OIDC_ARN}" >/dev/null 2>&1; then + echo " OIDC provider already exists, skipping" +else + echo " Creating OIDC provider..." + aws iam create-open-id-connect-provider \ + --url "https://token.actions.githubusercontent.com" \ + --client-id-list "sts.amazonaws.com" \ + --thumbprint-list "6938fd4d98bab03faadb97b34396831e3780aea1" "1c58a3a8518e8759bf075b76b750d4f2df264fcd" \ + --output text --query 'OpenIDConnectProviderArn' +fi + +# --- Trust Policy --- +SUB_CONDITIONS=$(printf '"%s"' "repo:${REPOS[0]}:ref:refs/heads/main") +for repo in "${REPOS[@]:1}"; do + SUB_CONDITIONS+=", \"repo:${repo}:ref:refs/heads/main\"" +done + +TRUST_POLICY=$(cat </dev/null 2>&1; then + echo " Role ${ROLE_NAME} exists, updating trust policy..." + aws iam update-assume-role-policy --role-name "${ROLE_NAME}" --policy-document "${TRUST_POLICY}" +else + echo " Creating role ${ROLE_NAME}..." + aws iam create-role \ + --role-name "${ROLE_NAME}" \ + --assume-role-policy-document "${TRUST_POLICY}" \ + --description "GitHub Actions role for bulk_executor e2e tests" \ + --output text --query 'Role.Arn' +fi + +# --- Inline Policy --- +echo " Attaching inline policy ${POLICY_NAME}..." +PERMISSIONS_POLICY=$(cat < Bootstrap complete" +echo " OIDC: ${OIDC_ARN}" +echo " Role: ${ROLE_ARN}" +echo " Repos: ${REPOS[*]}" +echo " Region: ${REGION}" +echo " Tables: ${READ_TABLE} (read), ${WRITE_TABLE} (write)" diff --git a/.github/scripts/e2e-switch-account.sh b/.github/scripts/e2e-switch-account.sh new file mode 100755 index 00000000..ab07dd09 --- /dev/null +++ b/.github/scripts/e2e-switch-account.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Switch e2e CI from one AWS account to another. +# Tears down the old account, bootstraps the new one. +# +# Prerequisites: +# - AWS CLI authenticated to the OLD account (for teardown) +# - You'll be prompted to switch credentials before bootstrap +# +# Usage: +# .github/scripts/e2e-switch-account.sh \ +# --old-account-id 654654401288 \ +# --new-account-id 111222333444 \ +# --region us-east-1 \ +# --read-table tiny-boat \ +# --write-table mini-boat \ +# --repos "awslabs/amazon-dynamodb-tools,relentlesscol/amazon-dynamodb-tools" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +usage() { + echo "Usage: $0 --old-account-id ID --new-account-id ID --region REGION --read-table TABLE --write-table TABLE --repos REPO1,REPO2" + exit 1 +} + +while [[ $# -gt 0 ]]; do + case $1 in + --old-account-id) OLD_ACCOUNT="$2"; shift 2 ;; + --new-account-id) NEW_ACCOUNT="$2"; shift 2 ;; + --region) REGION="$2"; shift 2 ;; + --read-table) READ_TABLE="$2"; shift 2 ;; + --write-table) WRITE_TABLE="$2"; shift 2 ;; + --repos) REPOS="$2"; shift 2 ;; + *) usage ;; + esac +done + +[[ -z "${OLD_ACCOUNT:-}" || -z "${NEW_ACCOUNT:-}" || -z "${REGION:-}" || -z "${READ_TABLE:-}" || -z "${WRITE_TABLE:-}" || -z "${REPOS:-}" ]] && usage + +echo "╔══════════════════════════════════════════════════╗" +echo "║ E2E Account Switch: ${OLD_ACCOUNT} → ${NEW_ACCOUNT} ║" +echo "╚══════════════════════════════════════════════════╝" +echo "" + +# --- Phase 1: Teardown old account --- +echo "── Phase 1: Teardown (account ${OLD_ACCOUNT}) ──" +echo "" +echo "Ensure AWS CLI is authenticated to ${OLD_ACCOUNT}." +read -rp "Press Enter to continue (or Ctrl+C to abort)..." +echo "" + +"${SCRIPT_DIR}/e2e-teardown.sh" --account-id "${OLD_ACCOUNT}" --repos "${REPOS}" + +echo "" + +# --- Phase 2: Bootstrap new account --- +echo "── Phase 2: Bootstrap (account ${NEW_ACCOUNT}) ──" +echo "" +echo "Switch AWS CLI credentials to ${NEW_ACCOUNT} now." +echo " e.g.: ada credentials update --account ${NEW_ACCOUNT} --provider isengard --role Admin --once" +echo "" +read -rp "Press Enter when ready (or Ctrl+C to abort)..." +echo "" + +"${SCRIPT_DIR}/e2e-bootstrap.sh" \ + --account-id "${NEW_ACCOUNT}" \ + --region "${REGION}" \ + --read-table "${READ_TABLE}" \ + --write-table "${WRITE_TABLE}" \ + --repos "${REPOS}" + +echo "" +echo "==> Account switch complete: ${OLD_ACCOUNT} → ${NEW_ACCOUNT}" diff --git a/.github/scripts/e2e-teardown.sh b/.github/scripts/e2e-teardown.sh new file mode 100755 index 00000000..e3b6bba0 --- /dev/null +++ b/.github/scripts/e2e-teardown.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Tear down GitHub Actions e2e infrastructure from a given AWS account. +# Removes: inline policy, IAM role, OIDC provider, GitHub secrets. +# +# Usage: +# .github/scripts/e2e-teardown.sh \ +# --account-id 123456789012 \ +# --repos "awslabs/amazon-dynamodb-tools,relentlesscol/amazon-dynamodb-tools" + +ROLE_NAME="github-actions-e2e-runner" +POLICY_NAME="e2e-test-access" + +usage() { + echo "Usage: $0 --account-id ID --repos REPO1,REPO2" + exit 1 +} + +while [[ $# -gt 0 ]]; do + case $1 in + --account-id) ACCOUNT_ID="$2"; shift 2 ;; + --repos) IFS=',' read -ra REPOS <<< "$2"; shift 2 ;; + *) usage ;; + esac +done + +[[ -z "${ACCOUNT_ID:-}" || ${#REPOS[@]} -eq 0 ]] && usage + +echo "==> Tearing down e2e CI from account ${ACCOUNT_ID}" + +# --- Inline Policy --- +if aws iam get-role-policy --role-name "${ROLE_NAME}" --policy-name "${POLICY_NAME}" >/dev/null 2>&1; then + echo " Deleting inline policy ${POLICY_NAME}..." + aws iam delete-role-policy --role-name "${ROLE_NAME}" --policy-name "${POLICY_NAME}" +else + echo " No inline policy found, skipping" +fi + +# --- IAM Role --- +if aws iam get-role --role-name "${ROLE_NAME}" >/dev/null 2>&1; then + echo " Deleting role ${ROLE_NAME}..." + aws iam delete-role --role-name "${ROLE_NAME}" +else + echo " No role found, skipping" +fi + +# --- OIDC Provider --- +OIDC_ARN="arn:aws:iam::${ACCOUNT_ID}:oidc-provider/token.actions.githubusercontent.com" +if aws iam get-open-id-connect-provider --open-id-connect-provider-arn "${OIDC_ARN}" >/dev/null 2>&1; then + echo " Deleting OIDC provider..." + aws iam delete-open-id-connect-provider --open-id-connect-provider-arn "${OIDC_ARN}" +else + echo " No OIDC provider found, skipping" +fi + +# --- GitHub Secrets --- +echo " Removing GitHub secrets..." +for repo in "${REPOS[@]}"; do + echo " ${repo}" + gh secret delete E2E_AWS_ROLE_ARN --repo "${repo}" 2>/dev/null || true + gh secret delete E2E_AWS_ACCOUNT_ID --repo "${repo}" 2>/dev/null || true + gh secret delete E2E_AWS_REGION --repo "${repo}" 2>/dev/null || true + gh secret delete E2E_READ_TABLE --repo "${repo}" 2>/dev/null || true + gh secret delete E2E_WRITE_TABLE --repo "${repo}" 2>/dev/null || true +done + +echo "" +echo "==> Teardown complete for account ${ACCOUNT_ID}" diff --git a/.github/workflows/bulk-executor-connector-smoke.yml b/.github/workflows/bulk-executor-connector-smoke.yml new file mode 100644 index 00000000..b1dd494c --- /dev/null +++ b/.github/workflows/bulk-executor-connector-smoke.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: connector smoke" + +on: + workflow_run: + workflows: ["bulk_executor: unit tests"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + connector-smoke: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run connector smoke + run: | + .venv/bin/pytest tests/e2e/connector/ -m e2e -v --tb=short --e2e-suite "connector smoke" diff --git a/.github/workflows/bulk-executor-e2e-copy.yml b/.github/workflows/bulk-executor-e2e-copy.yml new file mode 100644 index 00000000..cf1b029e --- /dev/null +++ b/.github/workflows/bulk-executor-e2e-copy.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: e2e copy" + +on: + workflow_run: + workflows: ["bulk_executor: connector smoke"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + e2e-copy: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run e2e copy + run: | + .venv/bin/pytest tests/e2e/commands/test_copy_smoke.py -m e2e -v --tb=short --e2e-suite "e2e copy" diff --git a/.github/workflows/bulk-executor-e2e-delete.yml b/.github/workflows/bulk-executor-e2e-delete.yml new file mode 100644 index 00000000..9d519a72 --- /dev/null +++ b/.github/workflows/bulk-executor-e2e-delete.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: e2e delete" + +on: + workflow_run: + workflows: ["bulk_executor: connector smoke"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + e2e-delete: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run e2e delete + run: | + .venv/bin/pytest tests/e2e/commands/test_delete_smoke.py -m e2e -v --tb=short --e2e-suite "e2e delete" diff --git a/.github/workflows/bulk-executor-e2e-diff.yml b/.github/workflows/bulk-executor-e2e-diff.yml new file mode 100644 index 00000000..93491415 --- /dev/null +++ b/.github/workflows/bulk-executor-e2e-diff.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: e2e diff" + +on: + workflow_run: + workflows: ["bulk_executor: connector smoke"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + e2e-diff: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run e2e diff + run: | + .venv/bin/pytest tests/e2e/commands/test_diff_smoke.py -m e2e -v --tb=short --e2e-suite "e2e diff" diff --git a/.github/workflows/bulk-executor-e2e-fill.yml b/.github/workflows/bulk-executor-e2e-fill.yml new file mode 100644 index 00000000..edffa51f --- /dev/null +++ b/.github/workflows/bulk-executor-e2e-fill.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: e2e fill" + +on: + workflow_run: + workflows: ["bulk_executor: connector smoke"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + e2e-fill: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run e2e fill + run: | + .venv/bin/pytest tests/e2e/commands/test_fill_smoke.py -m e2e -v --tb=short --e2e-suite "e2e fill" diff --git a/.github/workflows/bulk-executor-e2e-update.yml b/.github/workflows/bulk-executor-e2e-update.yml new file mode 100644 index 00000000..380fd163 --- /dev/null +++ b/.github/workflows/bulk-executor-e2e-update.yml @@ -0,0 +1,58 @@ +name: "bulk_executor: e2e update" + +on: + workflow_run: + workflows: ["bulk_executor: connector smoke"] + types: [completed] + branches: [main] + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + e2e-update: + runs-on: ubuntu-latest + if: > + (github.event_name == 'workflow_dispatch') || + (github.event.workflow_run.conclusion == 'success' && + (github.repository == 'awslabs/amazon-dynamodb-tools' || github.repository == 'relentlesscol/amazon-dynamodb-tools')) + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.E2E_AWS_ROLE_ARN }} + aws-region: ${{ secrets.E2E_AWS_REGION }} + + - name: Write e2e config + env: + E2E_ACCOUNT: ${{ secrets.E2E_AWS_ACCOUNT_ID }} + E2E_REGION: ${{ secrets.E2E_AWS_REGION }} + E2E_READ: ${{ secrets.E2E_READ_TABLE }} + E2E_WRITE: ${{ secrets.E2E_WRITE_TABLE }} + run: | + printf '{"aws_account_id":"%s","aws_region":"%s","read_table":"%s","write_table":"%s","bootstrap_confirmed":true}\n' \ + "$E2E_ACCOUNT" "$E2E_REGION" "$E2E_READ" "$E2E_WRITE" > tests/e2e/.e2e-config + + - name: Run e2e update + run: | + .venv/bin/pytest tests/e2e/commands/test_update_smoke.py -m e2e -v --tb=short --e2e-suite "e2e update" diff --git a/.github/workflows/bulk-executor-flake-retry.yml b/.github/workflows/bulk-executor-flake-retry.yml new file mode 100644 index 00000000..9de70e14 --- /dev/null +++ b/.github/workflows/bulk-executor-flake-retry.yml @@ -0,0 +1,64 @@ +name: "bulk_executor: flake retry" + +on: + workflow_run: + workflows: + - "bulk_executor: e2e fill" + - "bulk_executor: e2e copy" + - "bulk_executor: e2e update" + - "bulk_executor: e2e delete" + - "bulk_executor: e2e diff" + - "bulk_executor: connector smoke" + types: [completed] + branches: [main] + +permissions: + actions: write + contents: read + +jobs: + retry-on-failure: + runs-on: ubuntu-latest + if: github.event.workflow_run.conclusion == 'failure' + + steps: + - name: Check if this is already a retry + id: check + env: + GH_TOKEN: ${{ github.token }} + WORKFLOW_NAME: ${{ github.event.workflow_run.name }} + HEAD_SHA: ${{ github.event.workflow_run.head_sha }} + run: | + # Count recent failed runs for this workflow + commit + FAILED_COUNT=$(gh api \ + "repos/${{ github.repository }}/actions/runs?event=workflow_run&status=failure&head_sha=${HEAD_SHA}" \ + --jq "[.workflow_runs[] | select(.name == \"${WORKFLOW_NAME}\")] | length" \ + 2>/dev/null || echo "0") + + echo "failed_count=${FAILED_COUNT}" >> "$GITHUB_OUTPUT" + echo "Workflow '${WORKFLOW_NAME}' has ${FAILED_COUNT} failure(s) on commit ${HEAD_SHA:0:7}" + + - name: Retry if first failure + if: steps.check.outputs.failed_count == '1' + env: + GH_TOKEN: ${{ github.token }} + WORKFLOW_ID: ${{ github.event.workflow_run.workflow_id }} + run: | + echo "First failure — triggering retry..." + gh workflow run "${WORKFLOW_ID}" \ + --repo "${{ github.repository }}" \ + --ref "${{ github.event.workflow_run.head_branch }}" + + - name: Alert if repeated failure + if: steps.check.outputs.failed_count != '1' + env: + WORKFLOW_NAME: ${{ github.event.workflow_run.name }} + HEAD_SHA: ${{ github.event.workflow_run.head_sha }} + RUN_URL: ${{ github.event.workflow_run.html_url }} + run: | + echo "::error::REAL FAILURE (not a flake): '${WORKFLOW_NAME}' failed twice on ${HEAD_SHA:0:7}" + echo "" + echo "Run: ${RUN_URL}" + echo "" + echo "This is likely a genuine regression, not a transient Glue hiccup." + exit 1 diff --git a/.github/workflows/bulk-executor-unit-tests.yml b/.github/workflows/bulk-executor-unit-tests.yml new file mode 100644 index 00000000..0b2be413 --- /dev/null +++ b/.github/workflows/bulk-executor-unit-tests.yml @@ -0,0 +1,43 @@ +name: "bulk_executor: unit tests" + +on: + push: + branches: [main] + paths: + - 'tools/bulk_executor/**' + - '.github/workflows/bulk-executor-unit-tests.yml' + pull_request: + branches: [main] + paths: + - 'tools/bulk_executor/**' + - '.github/workflows/bulk-executor-unit-tests.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + unit-tests: + runs-on: ubuntu-latest + defaults: + run: + working-directory: tools/bulk_executor + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m venv .venv + .venv/bin/pip install --upgrade pip + .venv/bin/pip install -r requirements.txt + .venv/bin/pip install -r tests/requirements-test.txt + + - name: Run unit tests + run: | + .venv/bin/pytest tests/ --ignore=tests/e2e -v --tb=short --cov=server/src --cov=client/src --cov-branch --cov-report=term-missing From a99bcf1679cbabd1422f8f9f6757397675b4ccc4 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 17:59:29 -0700 Subject: [PATCH 02/22] fix(diff): skip writing empty diff segment files to S3 (bu-1a1) (#16) When using --s3 mode, diff_segment() was unconditionally writing a file per segment to S3 via put_object, even when no differences were found in that segment. This produced many empty 0-byte files cluttering the output bucket and adding unnecessary S3 API calls. Guard the put_object call with an emptiness check on the diff list. The return value (len(diff) == 0) is preserved so the aggregated count in run() still correctly reports "No differences found" when all segments are empty. Closes #183. --- tools/bulk_executor/server/src/python_modules/diff.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/bulk_executor/server/src/python_modules/diff.py b/tools/bulk_executor/server/src/python_modules/diff.py index f242ee51..979ebfb2 100644 --- a/tools/bulk_executor/server/src/python_modules/diff.py +++ b/tools/bulk_executor/server/src/python_modules/diff.py @@ -282,7 +282,8 @@ def diff_segment(stream_a_name, stream_b_name, monitor_options_a, monitor_option rate_limiter_worker_b.shutdown() if use_s3: - boto3.client('s3').put_object(Body="\n".join(diff), Bucket=bucket, Key=f"{job_id}/{segment}.txt") + if diff: + boto3.client('s3').put_object(Body="\n".join(diff), Bucket=bucket, Key=f"{job_id}/{segment}.txt") return len(diff) return diff[0:PRINT_LIMIT] From 18c02882fd9103eef1af72be29d460116677b06d Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 17:59:42 -0700 Subject: [PATCH 03/22] [bulk_executor] fix: exclude __pycache__ and dev cruft from python_modules.zip (#17) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Exclude __pycache__, .DS_Store, and dev cruft from python_modules.zip The module_zipper bundled everything under the source tree including bytecode caches, OS metadata files, and egg-info directories. Add exclusion lists for directories (__pycache__, .pytest_cache, .git, *.egg-info) and files (.DS_Store, Thumbs.db, .pyc, .pyo) that should never ship in the deployment zip. Fixes #174 --- .../client/src/utils/module_zipper.py | 22 +++ .../tests/client/utils/test_module_zipper.py | 151 ++++++++++++++++++ 2 files changed, 173 insertions(+) diff --git a/tools/bulk_executor/client/src/utils/module_zipper.py b/tools/bulk_executor/client/src/utils/module_zipper.py index 2052fe0d..9633ab58 100644 --- a/tools/bulk_executor/client/src/utils/module_zipper.py +++ b/tools/bulk_executor/client/src/utils/module_zipper.py @@ -8,6 +8,23 @@ ) from utils.logger import log +EXCLUDED_DIRS = {'__pycache__', '.pytest_cache', '.git', '*.egg-info'} +EXCLUDED_FILES = {'.DS_Store', 'Thumbs.db', '.gitignore'} +EXCLUDED_EXTENSIONS = {'.pyc', '.pyo'} + + +def _is_excluded_dir(name): + if name in EXCLUDED_DIRS: + return True + return name.endswith('.egg-info') + + +def _is_excluded_file(name): + if name in EXCLUDED_FILES: + return True + _, ext = os.path.splitext(name) + return ext in EXCLUDED_EXTENSIONS + def zip_module(): return _zip_module(PYTHON_MODULE_CLIENT_DIR_PATH, PYTHON_MODULE_CLIENT_ZIP_PATH) @@ -28,7 +45,12 @@ def _zip_module(source_path, zip_path): zipf.writestr(parent_dir + '/', '') for root, dirs, files in os.walk(source_path): + dirs[:] = [d for d in dirs if not _is_excluded_dir(d)] + for file in files: + if _is_excluded_file(file): + continue + file_path = os.path.join(root, file) # Skip any symlinked files out of an abundance of caution diff --git a/tools/bulk_executor/tests/client/utils/test_module_zipper.py b/tools/bulk_executor/tests/client/utils/test_module_zipper.py index 23bd6a7c..ff2edd09 100644 --- a/tools/bulk_executor/tests/client/utils/test_module_zipper.py +++ b/tools/bulk_executor/tests/client/utils/test_module_zipper.py @@ -5,6 +5,7 @@ infrastructure.constants - _zip_module(): happy path (writes parent dir entry, recurses os.walk, writes files, writes subdir entries to preserve empty dirs, skips symlinks), + exclusion of __pycache__, .DS_Store, .pyc/.pyo, and other junk, guard against zip_path inside source_path (raises ValueError → caught, returns False), exception during zipping (caught, log.error, returns False) @@ -149,6 +150,156 @@ def test_paths_normalized_to_absolute(self, tmp_path, monkeypatch): assert (tmp_path / 'out.zip').exists() +class TestZipModuleExclusions: + """Tests for file/directory exclusion logic.""" + + def test_pycache_excluded(self, tmp_path): + """__pycache__ directories and their contents are excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + cache = source / "__pycache__" + cache.mkdir() + (cache / "a.cpython-311.pyc").write_bytes(b"\x00") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert not any('__pycache__' in n for n in names) + + def test_nested_pycache_excluded(self, tmp_path): + """__pycache__ inside a subdirectory is also excluded.""" + source = tmp_path / "modules" + source.mkdir() + sub = source / "sub" + sub.mkdir() + (sub / "b.py").write_text("b") + cache = sub / "__pycache__" + cache.mkdir() + (cache / "b.cpython-311.pyc").write_bytes(b"\x00") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/sub/b.py' in names + assert not any('__pycache__' in n for n in names) + + def test_ds_store_excluded(self, tmp_path): + """.DS_Store files are excluded from the archive.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + (source / ".DS_Store").write_bytes(b"\x00\x00\x00\x01") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert not any('.DS_Store' in n for n in names) + + def test_pyc_files_excluded(self, tmp_path): + """.pyc files outside __pycache__ are also excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + (source / "a.pyc").write_bytes(b"\x00") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert 'modules/a.pyc' not in names + + def test_pyo_files_excluded(self, tmp_path): + """.pyo files are excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + (source / "a.pyo").write_bytes(b"\x00") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert 'modules/a.pyo' not in names + + def test_egg_info_dir_excluded(self, tmp_path): + """*.egg-info directories are excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + egg = source / "pkg.egg-info" + egg.mkdir() + (egg / "PKG-INFO").write_text("info") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert not any('egg-info' in n for n in names) + + def test_thumbs_db_excluded(self, tmp_path): + """Thumbs.db files are excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + (source / "Thumbs.db").write_bytes(b"\x00") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert 'modules/Thumbs.db' not in names + + def test_pytest_cache_excluded(self, tmp_path): + """.pytest_cache directories are excluded.""" + source = tmp_path / "modules" + source.mkdir() + (source / "a.py").write_text("a") + cache = source / ".pytest_cache" + cache.mkdir() + (cache / "CACHEDIR.TAG").write_text("tag") + + zip_path = tmp_path / "out.zip" + + with patch.object(module_zipper, 'log'): + assert module_zipper._zip_module(str(source), str(zip_path)) is True + + with zipfile.ZipFile(zip_path) as zf: + names = zf.namelist() + assert 'modules/a.py' in names + assert not any('.pytest_cache' in n for n in names) + + class TestZipModuleInternalGuard: """Tests for the zip-inside-source guard (lines 21-23).""" From ae37c49b0d493f6372e7bcc22669b038e0756a65 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 17:59:45 -0700 Subject: [PATCH 04/22] fix(bulk_executor): require PITR enabled for load command (bu-598) (#18) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The load command mutates table data but did not verify that Point-In-Time Recovery was enabled before proceeding. Other table-mutating commands (copy, update, fill) already pass pitr_enabled=True to validate_tables(), making this an oversight rather than a design choice. Without PITR, a failed or incorrect load has no recovery path — the user loses data with no automatic backup to restore from. Adding the flag ensures validate_tables() checks PITR status and exits with a clear error if it's disabled, matching the safety behavior of peer commands. Closes #179. --- tools/bulk_executor/client/src/python_modules/load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/bulk_executor/client/src/python_modules/load.py b/tools/bulk_executor/client/src/python_modules/load.py index 2b1e32bf..f26aa743 100644 --- a/tools/bulk_executor/client/src/python_modules/load.py +++ b/tools/bulk_executor/client/src/python_modules/load.py @@ -104,7 +104,7 @@ def run(env_configs): parser.error('--format should be "csv", "json" or "parquet"') result = args.__dict__ - utils.validate_tables(env_configs, parser, result['table']) + utils.validate_tables(env_configs, parser, result['table'], pitr_enabled=True) log.info(f"Running action '{result['verb']}' with arguments: {result}") From 32e186f5746e68f9dd07f57af19683c7e835964d Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 17:59:48 -0700 Subject: [PATCH 05/22] [bulk_executor] feat: warn if custom --XRole lacks minimum permissions (#19) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Warn if custom --XRole lacks minimum required permissions (#82) When a custom role is specified via --XRole, use simulate_principal_policy to check a representative set of required actions (DynamoDB, S3, CloudWatch, pricing, service quotas). Emit a warning listing any denied actions without blocking execution — the user may have grants we cannot detect (inline policies, resource-based policies, etc.). --- .../client/src/infrastructure/bootstrap.py | 34 ++++++++++ .../tests/client/test_bootstrap.py | 68 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/tools/bulk_executor/client/src/infrastructure/bootstrap.py b/tools/bulk_executor/client/src/infrastructure/bootstrap.py index 37014b01..b729b9f3 100644 --- a/tools/bulk_executor/client/src/infrastructure/bootstrap.py +++ b/tools/bulk_executor/client/src/infrastructure/bootstrap.py @@ -67,6 +67,7 @@ def _get_role_name(self, args): if not self._is_existing_role(role_param): print(f"Provided --XRole '{role_param}' name does not exist!") exit(1) + self._check_custom_role_permissions(role_param) return role_param # Handle standard role types @@ -182,6 +183,39 @@ def _is_existing_role(self, role_name): log.error(f'Unexpected error when checking for existing IAM Role: {e}') exit(1) + # Minimum actions a custom Glue execution role must be able to perform. + REQUIRED_ACTIONS = [ + 'dynamodb:DescribeTable', + 'dynamodb:Scan', + 's3:GetObject', + 's3:PutObject', + 'logs:CreateLogGroup', + 'logs:PutLogEvents', + 'pricing:GetProducts', + 'servicequotas:GetServiceQuota', + ] + + def _check_custom_role_permissions(self, role_name): + role_arn = f"arn:aws:iam::{self.aws_account_id}:role/{role_name}" + try: + response = self.iam_client.simulate_principal_policy( + PolicySourceArn=role_arn, + ActionNames=self.REQUIRED_ACTIONS, + ) + denied = [ + r['EvalActionName'] + for r in response.get('EvaluationResults', []) + if r.get('EvalDecision') != 'allowed' + ] + if denied: + log.warning( + f"Custom role '{role_name}' may be missing required permissions: " + f"{', '.join(denied)}. " + f"Execution may fail. See documentation for minimum required permissions." + ) + except Exception as e: + log.debug(f"Unable to verify permissions for role '{role_name}': {e}") + def _create_or_update_glue_job(self, args, is_create_allowed=True): glue_job_bucket = self._get_glue_job_bucket_name() diff --git a/tools/bulk_executor/tests/client/test_bootstrap.py b/tools/bulk_executor/tests/client/test_bootstrap.py index ce8d6761..060dda2f 100644 --- a/tools/bulk_executor/tests/client/test_bootstrap.py +++ b/tools/bulk_executor/tests/client/test_bootstrap.py @@ -151,6 +151,74 @@ def test_missing_custom_role_exits(self, bootstrap, capsys): assert 'NoSuchRole' in out +# -- _check_custom_role_permissions ------------------------------------- + +class TestCheckCustomRolePermissions: + """Coverage for the custom role minimum-permissions warning.""" + + def test_all_permissions_allowed_no_warning(self, bootstrap, caplog): + import logging + bootstrap.iam_client.simulate_principal_policy.return_value = { + 'EvaluationResults': [ + {'EvalActionName': action, 'EvalDecision': 'allowed'} + for action in bootstrap.REQUIRED_ACTIONS + ] + } + with caplog.at_level(logging.WARNING): + bootstrap._check_custom_role_permissions('GoodRole') + assert 'missing required permissions' not in caplog.text + + def test_some_permissions_denied_emits_warning(self, bootstrap, caplog): + import logging + results = [ + {'EvalActionName': 'dynamodb:DescribeTable', 'EvalDecision': 'allowed'}, + {'EvalActionName': 'dynamodb:Scan', 'EvalDecision': 'allowed'}, + {'EvalActionName': 's3:GetObject', 'EvalDecision': 'implicitDeny'}, + {'EvalActionName': 's3:PutObject', 'EvalDecision': 'allowed'}, + {'EvalActionName': 'logs:CreateLogGroup', 'EvalDecision': 'allowed'}, + {'EvalActionName': 'logs:PutLogEvents', 'EvalDecision': 'allowed'}, + {'EvalActionName': 'pricing:GetProducts', 'EvalDecision': 'explicitDeny'}, + {'EvalActionName': 'servicequotas:GetServiceQuota', 'EvalDecision': 'allowed'}, + ] + bootstrap.iam_client.simulate_principal_policy.return_value = { + 'EvaluationResults': results + } + with caplog.at_level(logging.WARNING): + bootstrap._check_custom_role_permissions('PartialRole') + assert 'missing required permissions' in caplog.text + assert 's3:GetObject' in caplog.text + assert 'pricing:GetProducts' in caplog.text + assert 'dynamodb:DescribeTable' not in caplog.text + + def test_simulate_api_error_does_not_block(self, bootstrap, caplog): + import logging + bootstrap.iam_client.simulate_principal_policy.side_effect = RuntimeError('access denied') + with caplog.at_level(logging.DEBUG): + bootstrap._check_custom_role_permissions('AnyRole') + assert 'Unable to verify permissions' in caplog.text + + def test_constructs_correct_arn(self, bootstrap): + bootstrap.iam_client.simulate_principal_policy.return_value = { + 'EvaluationResults': [] + } + bootstrap._check_custom_role_permissions('MyRole') + call_kwargs = bootstrap.iam_client.simulate_principal_policy.call_args.kwargs + assert call_kwargs['PolicySourceArn'] == 'arn:aws:iam::123456789012:role/MyRole' + assert call_kwargs['ActionNames'] == bootstrap.REQUIRED_ACTIONS + + def test_get_role_name_calls_check_for_custom_role(self, bootstrap): + bootstrap._is_existing_role = MagicMock(return_value=True) + bootstrap._check_custom_role_permissions = MagicMock() + bootstrap._get_role_name({'XRole': 'CustomRole'}) + bootstrap._check_custom_role_permissions.assert_called_once_with('CustomRole') + + def test_get_role_name_skips_check_for_standard_roles(self, bootstrap): + from infrastructure.constants import ROLE_TYPE_READ_ONLY + bootstrap._check_custom_role_permissions = MagicMock() + bootstrap._get_role_name({'XRole': ROLE_TYPE_READ_ONLY}) + bootstrap._check_custom_role_permissions.assert_not_called() + + # -- _is_existing_role -------------------------------------------------- class TestIsExistingRole: From bed6be5665dc2fa47819518ff30cc2953955be9c Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 17:59:51 -0700 Subject: [PATCH 06/22] [bulk_executor] feat: early-exit on non-recoverable systemic errors (#20) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * feat: add poison-pill early-exit for non-recoverable systemic errors (bu-i83b) When a Spark worker encounters a non-recoverable systemic error (AccessDeniedException, ValidationException, ExpiredTokenException, ResourceNotFoundException, ModuleNotFoundError, OutOfMemoryError), it writes a poison-pill marker to S3. All other workers check for this marker between scan pages and abort early instead of continuing to scan the entire table — which previously wasted minutes of compute and DynamoDB read capacity against a job that was guaranteed to fail. The mechanism uses the existing S3 bucket (already available for rate limiting) with a per-job-run key under server/poison-pill/. Workers rate-limit the HEAD check to once per 5 seconds to avoid S3 cost. The driver cleans up the marker in its finally block. The poison_pill_config parameter defaults to None in worker functions so existing test call-sites remain backwards-compatible. When None, a no-op implementation is used that never checks or signals. Affected verbs: copy, update, scancount, find (delete path). --- .../server/src/python_modules/copy.py | 19 +- .../server/src/python_modules/find.py | 15 +- .../src/python_modules/scancount/__init__.py | 16 +- .../src/python_modules/shared/poison_pill.py | 118 +++++++++++ .../src/python_modules/update/__init__.py | 19 +- tools/bulk_executor/tests/server/conftest.py | 12 +- .../tests/server/test_poison_pill.py | 200 ++++++++++++++++++ 7 files changed, 386 insertions(+), 13 deletions(-) create mode 100644 tools/bulk_executor/server/src/python_modules/shared/poison_pill.py create mode 100644 tools/bulk_executor/tests/server/test_poison_pill.py diff --git a/tools/bulk_executor/server/src/python_modules/copy.py b/tools/bulk_executor/server/src/python_modules/copy.py index a6d00869..d614abaf 100644 --- a/tools/bulk_executor/server/src/python_modules/copy.py +++ b/tools/bulk_executor/server/src/python_modules/copy.py @@ -7,6 +7,7 @@ sys.path.append('/server/src') from python_modules.shared.errors import get_error_message +from python_modules.shared.poison_pill import PoisonPillConfig, PoisonPillDriver, PoisonPillWorker, _NOOP as _NOOP_POISON_PILL from python_modules.shared.table_info import ( get_and_print_dynamodb_table_info, get_and_print_table_scan_cost, @@ -16,7 +17,7 @@ ) from python_modules.shared.rate_limiter import ( - RateLimiterAggregator, + RateLimiterAggregator, RateLimiterSharedConfig, RateLimiterWorker ) @@ -72,24 +73,30 @@ def run(job, spark_context, glue_context, parsed_args): # Since each task might generate errors, let's accumulate them and report intelligently error_accumulator = spark_context.accumulator([], ListAccumulator()) + poison_pill_config = PoisonPillConfig(bucket=bucket_name, job_run_id=job_run_id) + poison_pill_driver = PoisonPillDriver(poison_pill_config) + # Distribute work among partitions, each knowing what segment it's to handle try: parallelize_count = 400 rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) - rdd.foreach(lambda worker_id: _copy_data(source_table, target_table, source_monitor_options, target_monitor_options, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, source_rate_limiter_shared_config, target_rate_limiter_shared_config)) + rdd.foreach(lambda worker_id: _copy_data(source_table, target_table, source_monitor_options, target_monitor_options, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, source_rate_limiter_shared_config, target_rate_limiter_shared_config, poison_pill_config)) #rdd.count() except Exception as e: raise Exception(f"Error in parallel execution: {get_error_message(e)}") from None finally: source_rate_limiter_aggregator.shutdown() target_rate_limiter_aggregator.shutdown() + poison_pill_driver.cleanup() if error_accumulator.value: first_error = error_accumulator.value[0] raise Exception(first_error) from None print(f"Total records copied: {total_matched_accumulator.value:,}") -def _copy_data(source_table, target_table, source_monitor_options, target_monitor_options, segment, total_segments, total_matched_accumulator, error_accumulator, source_rate_limiter_shared_config, target_rate_limiter_shared_config): +def _copy_data(source_table, target_table, source_monitor_options, target_monitor_options, segment, total_segments, total_matched_accumulator, error_accumulator, source_rate_limiter_shared_config, target_rate_limiter_shared_config, poison_pill_config=None): + + poison_pill = PoisonPillWorker(poison_pill_config) if poison_pill_config else _NOOP_POISON_PILL # Let's hit the gas harder for this verb, at least for now XXX source_rl = RateLimiterWorker( @@ -128,6 +135,9 @@ def _copy_data(source_table, target_table, source_monitor_options, target_monito try: with dst.batch_writer() as batch: while True: + if poison_pill.check(): + break + resp = src.scan(**scan_kwargs) items = resp.get("Items", []) @@ -142,7 +152,8 @@ def _copy_data(source_table, target_table, source_monitor_options, target_monito scan_kwargs["ExclusiveStartKey"] = lek except Exception as e: error_accumulator.add([f"Error in worker {segment}: {get_error_message(e)}"]) - # Let control drop down to exit + if PoisonPillWorker.is_systemic_error(e): + poison_pill.signal(f"Worker {segment}: {get_error_message(e)}") finally: source_rl.shutdown() target_rl.shutdown() diff --git a/tools/bulk_executor/server/src/python_modules/find.py b/tools/bulk_executor/server/src/python_modules/find.py index 020f832d..4ca2a2a3 100644 --- a/tools/bulk_executor/server/src/python_modules/find.py +++ b/tools/bulk_executor/server/src/python_modules/find.py @@ -13,6 +13,7 @@ # Custom Library Imports sys.path.append('/server/src') from python_modules.shared.errors import * +from python_modules.shared.poison_pill import PoisonPillConfig, PoisonPillDriver, PoisonPillWorker, _NOOP as _NOOP_POISON_PILL from python_modules.shared.pricing import PricingUtility from python_modules.shared.rate_limiter import ( RateLimiterAggregator, @@ -180,7 +181,8 @@ def get_table_keys(table_name): elif DO_DELETE: keys = get_table_keys(DYNAMO_DB_TABLE_NAME) - def delete_partition(monitor_options, partition, shared_config): + def delete_partition(monitor_options, partition, shared_config, pp_config): + poison_pill = PoisonPillWorker(pp_config) rate_limiter_worker = RateLimiterWorker( shared_config=rate_limiter_shared_config, **monitor_options @@ -200,11 +202,16 @@ def delete_partition(monitor_options, partition, shared_config): try: with table.batch_writer() as batch: for record in partition: + if poison_pill.check(): + break try: item = json.loads(record) key = {k: item[k] for k in keys} batch.delete_item(Key=key) except Exception as e: + if PoisonPillWorker.is_systemic_error(e): + poison_pill.signal(f"Delete partition: {get_error_message(e)}") + raise print(f"Error deleting item {item}: {e}") finally: rate_limiter_worker.shutdown() @@ -228,15 +235,19 @@ def delete_partition(monitor_options, partition, shared_config): job_run_id=job_run_id ) + poison_pill_config = PoisonPillConfig(bucket=bucket_name, job_run_id=job_run_id) + poison_pill_driver = PoisonPillDriver(poison_pill_config) + rate_limiter_aggregator = RateLimiterAggregator(shared_config=rate_limiter_shared_config) monitor_options = get_dynamodb_throughput_configs(parsed_args, DYNAMO_DB_TABLE_NAME, modes=["write"], format="monitor") try: records.toJSON().foreachPartition( - lambda partition: delete_partition(monitor_options, partition, rate_limiter_shared_config) + lambda partition: delete_partition(monitor_options, partition, rate_limiter_shared_config, poison_pill_config) ) finally: rate_limiter_aggregator.shutdown() + poison_pill_driver.cleanup() print(f"Deleted {count:,} items") else: diff --git a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py index eb437c27..5d07d128 100644 --- a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py @@ -23,6 +23,7 @@ def decode(self, s): sys.path.append('/server/src') from python_modules.shared.errors import * from python_modules.shared.logger import log +from python_modules.shared.poison_pill import PoisonPillConfig, PoisonPillDriver, PoisonPillWorker, _NOOP as _NOOP_POISON_PILL from python_modules.shared.pricing import PricingUtility from python_modules.shared.rate_limiter import ( RateLimiterAggregator, @@ -77,16 +78,20 @@ def run(job, spark_context, glue_context, parsed_args): # Since each task might generate errors, let's accumulate them and report intelligently error_accumulator = spark_context.accumulator([], ListAccumulator()) + poison_pill_config = PoisonPillConfig(bucket=bucket_name, job_run_id=job_run_id) + poison_pill_driver = PoisonPillDriver(poison_pill_config) + # Distribute work among partitions, each knowing what segment it's to handle try: parallelize_count = 200 rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) - rdd.foreach(lambda worker_id: _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, rate_limiter_shared_config)) + rdd.foreach(lambda worker_id: _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config)) rdd.count() except Exception as e: raise Exception(f"Error in parallel execution: {get_error_message(e)}") from None finally: rate_limiter_aggregator.shutdown() + poison_pill_driver.cleanup() if error_accumulator.value: first_error = error_accumulator.value[0] raise Exception(first_error) from None @@ -94,7 +99,8 @@ def run(job, spark_context, glue_context, parsed_args): # Print the total records inserted using the accumulator after all tasks complete print(f"Total records counted: {total_matched_accumulator.value:,}") -def _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, segment, total_segments, total_matched_accumulator, error_accumulator, rate_limiter_shared_config): +def _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, segment, total_segments, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config=None): + poison_pill = PoisonPillWorker(poison_pill_config) if poison_pill_config else _NOOP_POISON_PILL rate_limiter_worker = RateLimiterWorker( shared_config=rate_limiter_shared_config, @@ -132,6 +138,9 @@ def _count_data(monitor_options, table_name, index_name, filter_expression, expr scan_kwargs["ExpressionAttributeValues"] = json.loads(expression_values, cls=DecimalEncoder) while True: + if poison_pill.check(): + break + response = table.scan(**scan_kwargs) # We do 50 retries within the SDK so shouldn't see a throttle response local_count += response.get("Count", 0) if "LastEvaluatedKey" not in response: @@ -139,7 +148,8 @@ def _count_data(monitor_options, table_name, index_name, filter_expression, expr scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"] except Exception as e: error_accumulator.add([f"Error in worker {segment}: {get_error_message(e)}"]) - # Let control drop down to exit + if PoisonPillWorker.is_systemic_error(e): + poison_pill.signal(f"Worker {segment}: {get_error_message(e)}") finally: rate_limiter_worker.shutdown() diff --git a/tools/bulk_executor/server/src/python_modules/shared/poison_pill.py b/tools/bulk_executor/server/src/python_modules/shared/poison_pill.py new file mode 100644 index 00000000..35711299 --- /dev/null +++ b/tools/bulk_executor/server/src/python_modules/shared/poison_pill.py @@ -0,0 +1,118 @@ +import time + +from boto3 import Session + +from .logger import log + +_POISON_KEY_SUFFIX = "poison-pill" +_CHECK_INTERVAL_SECONDS = 5 + + +class _NoOpPoisonPill: + """Placeholder when no poison-pill config is provided. All operations are no-ops.""" + + def signal(self, reason): + pass + + def check(self): + return False + + +_NOOP = _NoOpPoisonPill() + + +class PoisonPillConfig: + """Shared configuration for poison-pill coordination between driver and workers.""" + + def __init__(self, bucket, job_run_id): + self.bucket = bucket + self.key = f"server/poison-pill/{job_run_id}/{_POISON_KEY_SUFFIX}" + + +class PoisonPillDriver: + """Driver-side poison-pill lifecycle: cleanup on shutdown.""" + + def __init__(self, config): + self._config = config + self._s3 = Session().client("s3") + + def cleanup(self): + try: + self._s3.delete_object(Bucket=self._config.bucket, Key=self._config.key) + except Exception: + pass + + +class PoisonPillWorker: + """ + Worker-side poison-pill: signal fatal errors and check for abort. + + Call `check()` between scan pages. It rate-limits S3 reads to at most + once per _CHECK_INTERVAL_SECONDS. Call `signal(reason)` when this worker + hits a non-recoverable systemic error. + """ + + SYSTEMIC_ERRORS = frozenset([ + "AccessDeniedException", + "ModuleNotFoundError", + "OutOfMemoryError", + "ValidationException", + "ResourceNotFoundException", + "ExpiredTokenException", + ]) + + def __init__(self, config): + self._config = config + self._s3 = Session().client("s3") + self._last_check = 0.0 + self._poisoned = False + + def signal(self, reason): + """Write the poison marker so all other workers abort.""" + try: + self._s3.put_object( + Bucket=self._config.bucket, + Key=self._config.key, + Body=reason.encode("utf-8"), + ) + except Exception: + pass + self._poisoned = True + + def check(self): + """ + Return True if the job has been poisoned (another worker signaled abort). + + Rate-limits the S3 HEAD call to once per _CHECK_INTERVAL_SECONDS. + """ + if self._poisoned: + return True + + now = time.monotonic() + if now - self._last_check < _CHECK_INTERVAL_SECONDS: + return False + + self._last_check = now + try: + self._s3.head_object(Bucket=self._config.bucket, Key=self._config.key) + self._poisoned = True + return True + except self._s3.exceptions.NoSuchKey: + return False + except Exception: + return False + + @classmethod + def is_systemic_error(cls, error): + """Determine if an exception represents a non-recoverable systemic error.""" + error_str = str(type(error).__name__) + if error_str in cls.SYSTEMIC_ERRORS: + return True + + if hasattr(error, "response") and error.response: + error_code = error.response.get("Error", {}).get("Code", "") + if error_code in cls.SYSTEMIC_ERRORS: + return True + + msg = str(error) + return any(key in msg for key in cls.SYSTEMIC_ERRORS) diff --git a/tools/bulk_executor/server/src/python_modules/update/__init__.py b/tools/bulk_executor/server/src/python_modules/update/__init__.py index e68ea974..f980c8bb 100644 --- a/tools/bulk_executor/server/src/python_modules/update/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/update/__init__.py @@ -15,6 +15,7 @@ sys.path.append('/server/src') from python_modules.shared.errors import * from python_modules.shared.logger import log +from python_modules.shared.poison_pill import PoisonPillConfig, PoisonPillDriver, PoisonPillWorker, _NOOP as _NOOP_POISON_PILL from python_modules.shared.pricing import PricingUtility from python_modules.shared.rate_limiter import ( RateLimiterAggregator, @@ -75,15 +76,19 @@ def run(job, spark_context, glue_context, parsed_args): # Since each task might generate errors, let's accumulate them and report intelligently error_accumulator = spark_context.accumulator([], ListAccumulator()) + poison_pill_config = PoisonPillConfig(bucket=bucket_name, job_run_id=job_run_id) + poison_pill_driver = PoisonPillDriver(poison_pill_config) + # Distribute work among partitions, each knowing what segment it's to handle try: parallelize_count = 800 rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) - rdd.map(lambda worker_id: _update_data(monitor_options, table_name, generate, worker_id, parallelize_count, updated_accumulator, skipped_accumulator, failed_accumulator, error_accumulator, rate_limiter_shared_config)).collect() + rdd.map(lambda worker_id: _update_data(monitor_options, table_name, generate, worker_id, parallelize_count, updated_accumulator, skipped_accumulator, failed_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config)).collect() except Exception as e: raise Exception(f"Error in parallel execution: {get_error_message(e)}") from None finally: rate_limiter_aggregator.shutdown() + poison_pill_driver.cleanup() if error_accumulator.value: first_error = error_accumulator.value[0] raise Exception(first_error) from None @@ -93,7 +98,9 @@ def run(job, spark_context, glue_context, parsed_args): total = updated_accumulator.value + skipped_accumulator.value + failed_accumulator.value print(f"Processed {total:,} records: ({updated_accumulator.value:,} updates, {skipped_accumulator.value:,} non-updates, {failed_accumulator.value:,} conditions failed)") -def _update_data(monitor_options, table_name, generate, segment, total_segments, updated_accumulator, skipped_accumulator, failed_accumulator, error_accumulator, rate_limiter_shared_config): +def _update_data(monitor_options, table_name, generate, segment, total_segments, updated_accumulator, skipped_accumulator, failed_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config=None): + poison_pill = PoisonPillWorker(poison_pill_config) if poison_pill_config else _NOOP_POISON_PILL + rate_limiter_worker = RateLimiterWorker( shared_config=rate_limiter_shared_config, **monitor_options @@ -122,6 +129,9 @@ def _update_data(monitor_options, table_name, generate, segment, total_segments, try: while True: + if poison_pill.check(): + break + response = table.scan(**scan_kwargs) for item in response.get("Items", []): try: @@ -135,8 +145,10 @@ def _update_data(monitor_options, table_name, generate, segment, total_segments, except botocore.exceptions.ClientError as e: error_code = get_error_code(e) if error_code == DYNAMO_DB_THROTTLE_EXCEPTION: + poison_pill.signal(f"Worker {segment}: Throttling observed despite massive retries") exit("Throttling observed despite massive retries") elif error_code == DYNAMO_DB_VALIDATION_EXCEPTION: + poison_pill.signal(f"Worker {segment}: {get_error_message(e)}") exit(f"Validation exception (usually caused by the generator producing items incompatible with the table schema): {get_error_message(e)}") elif error_code == DYNAMO_DB_CONDITIONAL_CHECK_FAILED: print(f"UpdateItem condition expression failed, skipping... with kwargs: {update_kwargs}") @@ -151,7 +163,8 @@ def _update_data(monitor_options, table_name, generate, segment, total_segments, except Exception as e: error_accumulator.add([f"Error in worker {segment}: {get_error_message(e)}"]) - # Let control drop down to exit + if PoisonPillWorker.is_systemic_error(e): + poison_pill.signal(f"Worker {segment}: {get_error_message(e)}") finally: rate_limiter_worker.shutdown() diff --git a/tools/bulk_executor/tests/server/conftest.py b/tools/bulk_executor/tests/server/conftest.py index 5741af2c..f18c5379 100644 --- a/tools/bulk_executor/tests/server/conftest.py +++ b/tools/bulk_executor/tests/server/conftest.py @@ -107,7 +107,7 @@ def _write_dynamodb_stub(*args, **kwargs): sys.modules[f'{prefix}.glue_connector'].count_dynamodb_table = _count_dynamodb_stub sys.modules[f'{prefix}.glue_connector'].write_dynamodb_dataframe = _write_dynamodb_stub -# Import the real module — no pyspark dependency, so no mocking needed +# Import real modules — no pyspark dependency, so no mocking needed import importlib.util _spec = importlib.util.spec_from_file_location( "python_modules.shared.bulk_executor_error", @@ -118,6 +118,16 @@ def _write_dynamodb_stub(*args, **kwargs): for prefix in ['shared', 'python_modules.shared']: sys.modules[f'{prefix}.bulk_executor_error'] = _be_module +_pp_spec = importlib.util.spec_from_file_location( + "python_modules.shared.poison_pill", + str(__import__('pathlib').Path(__file__).resolve().parents[2] / "server/src/python_modules/shared/poison_pill.py") +) +_pp_module = importlib.util.module_from_spec(_pp_spec) +_pp_spec.loader.exec_module(_pp_module) +for prefix in ['shared', 'python_modules.shared']: + sys.modules[f'{prefix}.poison_pill'] = _pp_module + sys.modules[prefix].poison_pill = _pp_module + class MockRateLimiterWorker: def __init__(self, *args, **kwargs): diff --git a/tools/bulk_executor/tests/server/test_poison_pill.py b/tools/bulk_executor/tests/server/test_poison_pill.py new file mode 100644 index 00000000..b4745dfe --- /dev/null +++ b/tools/bulk_executor/tests/server/test_poison_pill.py @@ -0,0 +1,200 @@ +"""Unit tests for python_modules/shared/poison_pill.py. + +Covers: +- PoisonPillConfig: key derivation from bucket + job_run_id +- PoisonPillDriver: cleanup deletes the S3 marker, swallows errors +- PoisonPillWorker: + - signal() writes the poison marker to S3 + - check() returns False initially, True after signal + - check() rate-limits S3 HEAD calls + - is_systemic_error() classifies known fatal errors correctly +""" + +import sys +import time +from unittest.mock import MagicMock, patch + +import botocore.exceptions +import pytest + +import python_modules.shared.poison_pill as poison_pill_module +from python_modules.shared.poison_pill import ( + PoisonPillConfig, + PoisonPillDriver, + PoisonPillWorker, + _CHECK_INTERVAL_SECONDS, +) + + +class TestPoisonPillConfig: + def test_key_derivation(self): + config = PoisonPillConfig(bucket="my-bucket", job_run_id="jr_12345") + assert config.bucket == "my-bucket" + assert config.key == "server/poison-pill/jr_12345/poison-pill" + + def test_different_job_run_ids_produce_different_keys(self): + c1 = PoisonPillConfig(bucket="b", job_run_id="run-a") + c2 = PoisonPillConfig(bucket="b", job_run_id="run-b") + assert c1.key != c2.key + + +class TestPoisonPillDriver: + @patch.object(poison_pill_module, "Session") + def test_cleanup_deletes_object(self, mock_session_cls): + mock_s3 = MagicMock() + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + driver = PoisonPillDriver(config) + driver.cleanup() + + mock_s3.delete_object.assert_called_once_with(Bucket="bkt", Key=config.key) + + @patch.object(poison_pill_module, "Session") + def test_cleanup_swallows_errors(self, mock_session_cls): + mock_s3 = MagicMock() + mock_s3.delete_object.side_effect = Exception("network error") + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + driver = PoisonPillDriver(config) + driver.cleanup() # Should not raise + + +class TestPoisonPillWorkerSignal: + @patch.object(poison_pill_module, "Session") + def test_signal_writes_marker(self, mock_session_cls): + mock_s3 = MagicMock() + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker.signal("AccessDeniedException in worker 42") + + mock_s3.put_object.assert_called_once_with( + Bucket="bkt", + Key=config.key, + Body=b"AccessDeniedException in worker 42", + ) + + @patch.object(poison_pill_module, "Session") + def test_signal_sets_poisoned_flag(self, mock_session_cls): + mock_s3 = MagicMock() + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker.signal("fatal") + + assert worker.check() is True + + @patch.object(poison_pill_module, "Session") + def test_signal_swallows_s3_errors(self, mock_session_cls): + mock_s3 = MagicMock() + mock_s3.put_object.side_effect = Exception("timeout") + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker.signal("fatal") # Should not raise + + +class TestPoisonPillWorkerCheck: + @patch.object(poison_pill_module, "Session") + def test_check_returns_false_when_no_marker(self, mock_session_cls): + mock_s3 = MagicMock() + no_such_key = type("NoSuchKey", (Exception,), {}) + mock_s3.exceptions.NoSuchKey = no_such_key + mock_s3.head_object.side_effect = no_such_key() + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker._last_check = 0.0 # Force check + + assert worker.check() is False + + @patch.object(poison_pill_module, "Session") + def test_check_returns_true_when_marker_exists(self, mock_session_cls): + mock_s3 = MagicMock() + mock_s3.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + mock_s3.head_object.return_value = {} # Object exists + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker._last_check = 0.0 # Force check + + assert worker.check() is True + + @patch.object(poison_pill_module, "Session") + def test_check_rate_limits_calls(self, mock_session_cls): + mock_s3 = MagicMock() + no_such_key = type("NoSuchKey", (Exception,), {}) + mock_s3.exceptions.NoSuchKey = no_such_key + mock_s3.head_object.side_effect = no_such_key() + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker._last_check = time.monotonic() # Just checked + + # Should return False without calling S3 (rate limited) + assert worker.check() is False + mock_s3.head_object.assert_not_called() + + @patch.object(poison_pill_module, "Session") + def test_check_swallows_unexpected_errors(self, mock_session_cls): + mock_s3 = MagicMock() + mock_s3.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + mock_s3.head_object.side_effect = RuntimeError("unexpected") + mock_session_cls.return_value.client.return_value = mock_s3 + + config = PoisonPillConfig(bucket="bkt", job_run_id="jr1") + worker = PoisonPillWorker(config) + worker._last_check = 0.0 + + assert worker.check() is False # Should not raise + + +class TestIsSystemicError: + def test_access_denied_by_error_code(self): + e = botocore.exceptions.ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "denied"}}, + "Scan" + ) + assert PoisonPillWorker.is_systemic_error(e) is True + + def test_validation_exception_by_error_code(self): + e = botocore.exceptions.ClientError( + {"Error": {"Code": "ValidationException", "Message": "bad"}}, + "Scan" + ) + assert PoisonPillWorker.is_systemic_error(e) is True + + def test_throttle_is_not_systemic(self): + e = botocore.exceptions.ClientError( + {"Error": {"Code": "ProvisionedThroughputExceededException", "Message": "slow down"}}, + "Scan" + ) + # Throttling with retries IS systemic (it means retries are exhausted) + assert PoisonPillWorker.is_systemic_error(e) is False + + def test_module_not_found_by_class_name(self): + e = ModuleNotFoundError("No module named 'foo'") + assert PoisonPillWorker.is_systemic_error(e) is True + + def test_expired_token_in_message(self): + e = Exception("Something ExpiredTokenException something") + assert PoisonPillWorker.is_systemic_error(e) is True + + def test_conditional_check_failed_is_not_systemic(self): + e = botocore.exceptions.ClientError( + {"Error": {"Code": "ConditionalCheckFailedException", "Message": "condition"}}, + "UpdateItem" + ) + assert PoisonPillWorker.is_systemic_error(e) is False + + def test_generic_exception_is_not_systemic(self): + e = ValueError("some value error") + assert PoisonPillWorker.is_systemic_error(e) is False From 11c12da610bcf88377d39fdbd6a555dcd08ca838 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:00 -0700 Subject: [PATCH 07/22] fix(rate_limiter): downgrade noisy Initializing log to debug (#181) (#24) The log.info call in RateLimiterAggregator.__init__ was confusing because it printed at INFO level during normal operation. Downgrade to debug. Also fix the format string which was printing bucket for both Bucket and Prefix fields. Fixes #181 --- .../server/src/python_modules/shared/rate_limiter/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py b/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py index 65072f3a..b64a115b 100644 --- a/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py @@ -33,7 +33,7 @@ class RateLimiterAggregator: modes (none to many list of ("read", "write")): The expected execution modes of the DynamoDB actions requiring rate limiting. """ def __init__(self, shared_config): - log.info(f"Initializing...Bucket:{shared_config.bucket}, Prefix:{shared_config.bucket}") + log.debug(f"Initializing...Bucket:{shared_config.bucket}, Prefix:{shared_config.prefix}") self.rate_limiter_monitor_aggregator = DistributedDynamoDBMonitorAggregator( session=Session(), From 55c3ff9c567ef8d6fdd1bb8e10b12a5e5093aa23 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:03 -0700 Subject: [PATCH 08/22] feat(load): report write rate at load start (bu-p8y) (#25) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the load command begins writing, log the write rate (WCU/s) being used and whether it was explicitly set by the user via --XMaxWriteRate or automatically determined by the DynamoDB connector. This addresses issue #182 — users had no visibility into what write throughput the load operation would consume. The message appears after the cost estimate and before the actual write begins, so operators can abort if the rate is unexpected. --- .../server/src/python_modules/load/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/bulk_executor/server/src/python_modules/load/__init__.py b/tools/bulk_executor/server/src/python_modules/load/__init__.py index 91462e0d..095ed401 100644 --- a/tools/bulk_executor/server/src/python_modules/load/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/load/__init__.py @@ -104,6 +104,12 @@ def run(job, spark_context, glue_context, parsed_args): session = boto3.Session() print_dynamodb_table_info(session, table_name, count, check_dynamic_frame_avg_size(dynamicFrame)) + write_rate = parsed_args.get('XMaxWriteRate') + if write_rate is not None: + log.info(f"Write rate: {write_rate} WCU/s (explicitly set via --XMaxWriteRate)") + else: + log.info("Write rate: automatically determined by the DynamoDB connector") + df = dynamicFrame.repartition(30).toDF() write_dynamodb_dataframe( glue_context, df, table_name, parsed_args) From b2e44b79184ac3cd545113d43e1c94f9ce73616f Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:07 -0700 Subject: [PATCH 09/22] fix(find): preserve DynamoDB types in find-to-S3 output (bu-waj) (#26) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The find verb previously serialized records through an intermediate spark.read.json(records.toJSON()) step before writing to S3. This re-inferred the schema from JSON strings, losing type fidelity for DynamoDB complex types (maps, lists, sets, binary, boolean, null). Numbers could also lose precision during the re-inference pass. Replace the toJSON→read.json→write.json pipeline with a direct records.write.mode('overwrite').json(location) call, which serializes the DataFrame using its existing Spark schema — exactly the schema the Glue DynamoDB connector produced on read. This preserves all attribute types through the find→S3→load round trip. Remove the now-unused SparkSession import (it was only needed for the intermediate spark.read.json() call). Tested: 1252 unit tests pass (3 new tests validate the direct-write path and absence of the re-inference step). Closes #184 --- .../server/src/python_modules/find.py | 12 +--- tools/bulk_executor/tests/server/test_find.py | 70 +++++++++++++++++++ 2 files changed, 73 insertions(+), 9 deletions(-) diff --git a/tools/bulk_executor/server/src/python_modules/find.py b/tools/bulk_executor/server/src/python_modules/find.py index 4ca2a2a3..92cad96a 100644 --- a/tools/bulk_executor/server/src/python_modules/find.py +++ b/tools/bulk_executor/server/src/python_modules/find.py @@ -7,7 +7,6 @@ import boto3 from awsglue.transforms import * from botocore.config import Config -from pyspark.sql import SparkSession from pyspark.sql.functions import asc, desc # Custom Library Imports @@ -154,14 +153,9 @@ def get_table_keys(table_name): job_run_id = parsed_args.get("JOB_RUN_ID") s3_output_location = f"s3://{bucket_name}/output/{job_run_id}" - # With a --limit, this produces one file with a name like part-00000-8c460443-6d45-4d11-b9ef-0cd84c21a45a-c000.json - # because the limit moves all the data to a single worker - # Without a limit, this produces about 200 files with names similar to that - # Adding coalesce(10) gets us down to 10 files, but testing against a large table showed that slower - spark = SparkSession(spark_context) - json_rdd = records.toJSON() - json_df = spark.read.json(json_rdd) - json_df.write.mode("overwrite").json(s3_output_location) + # Write the DataFrame directly to S3 as JSON Lines, preserving the + # connector's schema (type fidelity for maps, numbers, booleans, etc.). + records.write.mode("overwrite").json(s3_output_location) # Print the top N many TOP_N = 10 diff --git a/tools/bulk_executor/tests/server/test_find.py b/tools/bulk_executor/tests/server/test_find.py index 46c88d95..f412f117 100644 --- a/tools/bulk_executor/tests/server/test_find.py +++ b/tools/bulk_executor/tests/server/test_find.py @@ -757,6 +757,76 @@ def test_find_writes_count_items_message( assert 'Wrote 42 items in JSON format' in out +# --- run(): DO_FIND — direct DataFrame write (no schema re-inference) -------- + +class TestRunFindDirectWrite: + """DO_FIND writes JSON to S3 via records.write.json() directly, + preserving the DynamoDB connector's schema without re-inferring + via spark.read.json(). Addresses issue #184.""" + + def test_find_writes_dataframe_directly_to_s3( + self, monkeypatch, table_info_mocks, boto3_session_mock, base_args + ): + """records.write.mode('overwrite').json(location) is called.""" + df = _pyspark_sql.MagicMock if hasattr(_pyspark_sql, 'MagicMock') else MagicMock() + df = MagicMock() + df.cache.return_value = df + df.count.return_value = 2 + df.limit.return_value.toJSON.return_value.collect.return_value = ['{"a":1}', '{"b":2}'] + writer = MagicMock() + df.write = writer + writer.mode.return_value = writer + + monkeypatch.setattr(find_module, 'read_dynamodb_dataframe', lambda *a, **kw: df) + + find_module.run(MagicMock(), MagicMock(), MagicMock(), base_args) + + writer.mode.assert_called_once_with("overwrite") + writer.json.assert_called_once_with("s3://my-bucket/output/run-123") + + def test_find_does_not_use_spark_read_json( + self, monkeypatch, table_info_mocks, boto3_session_mock, base_args + ): + """No intermediate spark.read.json() — schema is not re-inferred.""" + df = MagicMock() + df.cache.return_value = df + df.count.return_value = 1 + df.limit.return_value.toJSON.return_value.collect.return_value = ['{"x":1}'] + writer = MagicMock() + df.write = writer + writer.mode.return_value = writer + + monkeypatch.setattr(find_module, 'read_dynamodb_dataframe', lambda *a, **kw: df) + + spark_ctx = MagicMock() + find_module.run(MagicMock(), spark_ctx, MagicMock(), base_args) + + assert not hasattr(find_module, 'SparkSession'), \ + "SparkSession import removed — no re-inference path" + + def test_find_prints_top_n_and_count( + self, monkeypatch, table_info_mocks, boto3_session_mock, base_args, capsys + ): + """Top-N preview and written count still work after the write change.""" + df = MagicMock() + df.cache.return_value = df + df.count.return_value = 15 + records = [f'{{"id":{i}}}' for i in range(10)] + df.limit.return_value.toJSON.return_value.collect.return_value = records + writer = MagicMock() + df.write = writer + writer.mode.return_value = writer + + monkeypatch.setattr(find_module, 'read_dynamodb_dataframe', lambda *a, **kw: df) + + find_module.run(MagicMock(), MagicMock(), MagicMock(), base_args) + + out = capsys.readouterr().out + assert 'First 10 matching items:' in out + assert '5 more not printed' in out + assert 'Wrote 15 items in JSON format' in out + + # --- run(): DO_DELETE branch -------------------------------------------------- @pytest.mark.skip(reason="Asserts against legacy DynamicFrame code path; verb now goes through python_modules.shared.glue_connector wrapper. Followup: rewrite to assert against the wrapper boundary.") From 386bb71f776f75254ea0b406ad3a6db359e8524e Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:35 -0700 Subject: [PATCH 10/22] Remove redundant TableName from scan_kwargs in update module (#5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Remove redundant TableName from scan_kwargs in update module Table.scan() is called on a boto3 Table resource that already knows its table name. Passing TableName is dead code that could mask bugs if the variable and resource ever diverge. --- .../server/src/python_modules/update/__init__.py | 1 - tools/bulk_executor/tests/server/test_update.py | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tools/bulk_executor/server/src/python_modules/update/__init__.py b/tools/bulk_executor/server/src/python_modules/update/__init__.py index f980c8bb..3f0cdb2f 100644 --- a/tools/bulk_executor/server/src/python_modules/update/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/update/__init__.py @@ -122,7 +122,6 @@ def _update_data(monitor_options, table_name, generate, segment, total_segments, skipped_count = 0 failed_count = 0 scan_kwargs = { - "TableName": table_name, "Segment": segment, "TotalSegments": total_segments } diff --git a/tools/bulk_executor/tests/server/test_update.py b/tools/bulk_executor/tests/server/test_update.py index 6664f68e..26ec8029 100644 --- a/tools/bulk_executor/tests/server/test_update.py +++ b/tools/bulk_executor/tests/server/test_update.py @@ -559,7 +559,7 @@ def scan_capture(**kwargs): updated_acc.add.assert_called_once_with(3) def test_scan_kwargs_include_segment_and_total(self, monkeypatch): - """Lines 117-121: scan_kwargs includes TableName, Segment, TotalSegments.""" + """Lines 117-120: scan_kwargs includes Segment, TotalSegments (not TableName — Table resource knows its name).""" scan_kwargs_seen = [] table = MagicMock() @@ -576,7 +576,7 @@ def scan_capture(**kwargs): MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock() ) - assert scan_kwargs_seen[0]['TableName'] == 'my-tbl' + assert 'TableName' not in scan_kwargs_seen[0] assert scan_kwargs_seen[0]['Segment'] == 7 assert scan_kwargs_seen[0]['TotalSegments'] == 100 From 2560d5b13bddfa642868a5a3140354df6b75c6a0 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:38 -0700 Subject: [PATCH 11/22] Show clean error messages for bad parameters instead of stack traces (#6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Show clean error messages for bad parameters instead of stack traces Wrap the driver's main execution in try/except for ClientError and BotoCoreError so users see "Error: AccessDeniedException — " rather than a full Python traceback when e.g. they pass an S3 bucket they don't own or a table they can't access. Fixes #137 --- tools/bulk_executor/bulk | 43 +++-- .../client/test_bulk_driver_error_handling.py | 177 ++++++++++++++++++ 2 files changed, 205 insertions(+), 15 deletions(-) create mode 100644 tools/bulk_executor/tests/client/test_bulk_driver_error_handling.py diff --git a/tools/bulk_executor/bulk b/tools/bulk_executor/bulk index de06c432..fa0e242d 100755 --- a/tools/bulk_executor/bulk +++ b/tools/bulk_executor/bulk @@ -4,6 +4,7 @@ import importlib import sys sys.path.append('./client/src/') +import botocore.exceptions import utils from env_configs import EnvConfigs from infrastructure import BootstrapInfrastructure @@ -83,19 +84,31 @@ action_script_function = _get_action_script_function(action) if(action_script_function is None): exit(1) # No verb to take additional action on. -env_configs = _get_env_configs() +try: + env_configs = _get_env_configs() -# Run client side verb script -is_client_and_server_action, processed_args = action_script_function(env_configs) # Run the client side verb function - -if(is_client_and_server_action): # Only run the server if the client script DNE or verifies the server script should also be run. - # Run server side verb script - args = utils.get_args_from_processed_args(processed_args) - script_args = utils.convert_client_dict_to_script_args(processed_args) - - # Dev Mode - Push new code to S3 without a full bootstrap - if(args.get('XDev')): - BootstrapInfrastructure(env_configs).update_python_modules_in_s3() - - # Run the server side verb function (Glue Job) - BulkDynamoDbRunner(env_configs).run(args, script_args) + # Run client side verb script + is_client_and_server_action, processed_args = action_script_function(env_configs) # Run the client side verb function + + if(is_client_and_server_action): # Only run the server if the client script DNE or verifies the server script should also be run. + # Run server side verb script + args = utils.get_args_from_processed_args(processed_args) + script_args = utils.convert_client_dict_to_script_args(processed_args) + + # Dev Mode - Push new code to S3 without a full bootstrap + if(args.get('XDev')): + BootstrapInfrastructure(env_configs).update_python_modules_in_s3() + + # Run the server side verb function (Glue Job) + BulkDynamoDbRunner(env_configs).run(args, script_args) + +except botocore.exceptions.ClientError as e: + error_code = e.response['Error']['Code'] + error_message = e.response['Error']['Message'] + print(f"\nError: {error_code} — {error_message}", file=sys.stderr) + print("Job failed.", file=sys.stderr) + exit(1) +except botocore.exceptions.BotoCoreError as e: + print(f"\nError: {e}", file=sys.stderr) + print("Job failed.", file=sys.stderr) + exit(1) diff --git a/tools/bulk_executor/tests/client/test_bulk_driver_error_handling.py b/tools/bulk_executor/tests/client/test_bulk_driver_error_handling.py new file mode 100644 index 00000000..25d0382e --- /dev/null +++ b/tools/bulk_executor/tests/client/test_bulk_driver_error_handling.py @@ -0,0 +1,177 @@ +"""Tests for the top-level error handling in the `bulk` driver script. + +Verifies that known AWS errors (ClientError, BotoCoreError) produce a clean +one-line message + exit(1) instead of a full Python stack trace. + +The `bulk` script wraps its main execution in a try/except for ClientError +and BotoCoreError. These tests exercise that code path by running the script +as a subprocess from the correct working directory. +""" + +import subprocess +import sys +import os +import tempfile + +import pytest + +BULK_DIR = os.path.abspath( + os.path.join(os.path.dirname(__file__), '..', '..') +) + + +def _run_bulk_subprocess(inject_code): + """Run the bulk script with injected setup code. + + The wrapper adds necessary paths and patches, then exec's the bulk script. + """ + preamble = ( + "import sys, os\n" + f"os.chdir('{BULK_DIR}')\n" + f"sys.path.insert(0, '{BULK_DIR}')\n" + f"sys.path.insert(0, '{BULK_DIR}/client/src')\n" + ) + script = preamble + inject_code + "\nexec(compile(open('bulk').read(), 'bulk', 'exec'))\n" + + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: + f.write(script) + f.flush() + try: + result = subprocess.run( + [sys.executable, f.name], + capture_output=True, + text=True, + cwd=BULK_DIR, + timeout=10, + ) + finally: + os.unlink(f.name) + return result + + +INJECT_CLIENT_ERROR_ACCESS_DENIED = """\ +from unittest.mock import patch, MagicMock +import botocore.exceptions + +sys.argv = ['bulk', 'copy', '--source', 'src-tbl', '--target', 'dst-tbl'] + +error_response = { + 'Error': { + 'Code': 'AccessDeniedException', + 'Message': 'Not authorized to perform dynamodb:DescribeTable on table/src-tbl' + } +} + +mock_env = MagicMock(aws_region='us-east-1', aws_account_id='123456789012') + +def fake_action(env_configs): + raise botocore.exceptions.ClientError(error_response, 'DescribeTable') + +patch('env_configs.EnvConfigs', return_value=mock_env).start() +patch('utils.logger.init').start() +_m = patch('importlib.import_module').start() +_mod = MagicMock() +_mod.run = fake_action +_m.return_value = _mod +""" + +INJECT_CLIENT_ERROR_RESOURCE_NOT_FOUND = """\ +from unittest.mock import patch, MagicMock +import botocore.exceptions + +sys.argv = ['bulk', 'copy', '--source', 'src-tbl', '--target', 'dst-tbl'] + +error_response = { + 'Error': { + 'Code': 'ResourceNotFoundException', + 'Message': 'Requested resource not found: Table: nonexistent-table not found' + } +} + +mock_env = MagicMock(aws_region='us-east-1', aws_account_id='123456789012') + +def fake_action(env_configs): + raise botocore.exceptions.ClientError(error_response, 'DescribeTable') + +patch('env_configs.EnvConfigs', return_value=mock_env).start() +patch('utils.logger.init').start() +_m = patch('importlib.import_module').start() +_mod = MagicMock() +_mod.run = fake_action +_m.return_value = _mod +""" + +INJECT_BOTOCORE_ENDPOINT_ERROR = """\ +from unittest.mock import patch, MagicMock +import botocore.exceptions + +sys.argv = ['bulk', 'copy', '--source', 'src-tbl', '--target', 'dst-tbl'] + +mock_env = MagicMock(aws_region='us-east-1', aws_account_id='123456789012') + +def fake_action(env_configs): + raise botocore.exceptions.EndpointConnectionError( + endpoint_url='https://dynamodb.bad-region.amazonaws.com') + +patch('env_configs.EnvConfigs', return_value=mock_env).start() +patch('utils.logger.init').start() +_m = patch('importlib.import_module').start() +_mod = MagicMock() +_mod.run = fake_action +_m.return_value = _mod +""" + +INJECT_TYPE_ERROR = """\ +from unittest.mock import patch, MagicMock + +sys.argv = ['bulk', 'copy', '--source', 'src-tbl', '--target', 'dst-tbl'] + +mock_env = MagicMock(aws_region='us-east-1', aws_account_id='123456789012') + +def fake_action(env_configs): + raise TypeError("something internal broke") + +patch('env_configs.EnvConfigs', return_value=mock_env).start() +patch('utils.logger.init').start() +_m = patch('importlib.import_module').start() +_mod = MagicMock() +_mod.run = fake_action +_m.return_value = _mod +""" + + +class TestClientErrorHandling: + """Verify ClientError triggers clean message + exit(1), no stack trace.""" + + def test_access_denied_shows_code_and_message(self): + result = _run_bulk_subprocess(INJECT_CLIENT_ERROR_ACCESS_DENIED) + assert result.returncode == 1 + assert 'Traceback' not in result.stderr + assert 'AccessDeniedException' in result.stderr + assert 'Not authorized' in result.stderr + assert 'Job failed.' in result.stderr + + def test_resource_not_found_shows_clean_error(self): + result = _run_bulk_subprocess(INJECT_CLIENT_ERROR_RESOURCE_NOT_FOUND) + assert result.returncode == 1 + assert 'Traceback' not in result.stderr + assert 'ResourceNotFoundException' in result.stderr + assert 'Job failed.' in result.stderr + + +class TestBotoCoreErrorHandling: + """Verify BotoCoreError (non-ClientError) triggers clean message.""" + + def test_endpoint_connection_error_clean_output(self): + result = _run_bulk_subprocess(INJECT_BOTOCORE_ENDPOINT_ERROR) + assert result.returncode == 1 + assert 'Traceback' not in result.stderr + assert 'Job failed.' in result.stderr + + def test_non_aws_error_still_shows_traceback(self): + """Non-AWS errors (e.g. TypeError) should still produce a traceback + so bugs are visible during development.""" + result = _run_bulk_subprocess(INJECT_TYPE_ERROR) + assert result.returncode == 1 + assert 'Traceback' in result.stderr + assert 'TypeError' in result.stderr From 0e1ded6bd293dc92fe2170428a501a3c1e7512bd Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:42 -0700 Subject: [PATCH 12/22] Add dedicated unit tests for rate_limiter modules (#7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Add dedicated unit tests for rate_limiter modules TokenBucket (24 tests): refill, deduct, wait_until_positive, reconfigure, negative balance, concurrency safety. DynamoDBMonitor (22 tests): event hook registration, capacity tracking, bucket deduction, rate setters, reporting lifecycle. DistributedDynamoDBMonitorWorker (14 tests): S3 sync upload, rate computation from deltas, scaling logic from aggregator summary, cleanup. DistributedDynamoDBMonitorAggregator (14 tests): multi-worker aggregation, staleness cutoff, corrupt data handling, summary file exclusion, cleanup. These modules are the core throttling safety mechanism and previously had zero dedicated test coverage (all verb tests mock them out). --- .../tests/server/rate_limiter/__init__.py | 0 .../tests/server/rate_limiter/conftest.py | 44 +++ .../test_distributed_aggregator.py | 232 ++++++++++++++++ .../rate_limiter/test_distributed_worker.py | 254 ++++++++++++++++++ .../rate_limiter/test_dynamodb_monitor.py | 206 ++++++++++++++ .../server/rate_limiter/test_token_bucket.py | 218 +++++++++++++++ 6 files changed, 954 insertions(+) create mode 100644 tools/bulk_executor/tests/server/rate_limiter/__init__.py create mode 100644 tools/bulk_executor/tests/server/rate_limiter/conftest.py create mode 100644 tools/bulk_executor/tests/server/rate_limiter/test_distributed_aggregator.py create mode 100644 tools/bulk_executor/tests/server/rate_limiter/test_distributed_worker.py create mode 100644 tools/bulk_executor/tests/server/rate_limiter/test_dynamodb_monitor.py create mode 100644 tools/bulk_executor/tests/server/rate_limiter/test_token_bucket.py diff --git a/tools/bulk_executor/tests/server/rate_limiter/__init__.py b/tools/bulk_executor/tests/server/rate_limiter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/bulk_executor/tests/server/rate_limiter/conftest.py b/tools/bulk_executor/tests/server/rate_limiter/conftest.py new file mode 100644 index 00000000..97e3e3f1 --- /dev/null +++ b/tools/bulk_executor/tests/server/rate_limiter/conftest.py @@ -0,0 +1,44 @@ +"""Conftest for rate_limiter tests — imports real modules instead of mocks.""" +import sys +import types +import logging +import importlib.util +from pathlib import Path +from unittest.mock import Mock + +_shared_parent = Path(__file__).resolve().parents[3] / "server" / "src" / "python_modules" / "shared" +_rl_path = _shared_parent / "rate_limiter" + +# Set up logger mock (rate_limiter modules import from ..logger) +_real_logger = logging.getLogger('rate_limiter_tests') +_real_logger.setLevel(logging.DEBUG) + +_logger_module = types.ModuleType('python_modules.shared.logger') +_logger_module.log = _real_logger +_logger_module.init = Mock() + +# Overwrite the mocked entries with real namespace packages +_pm = types.ModuleType('python_modules') +_pm.__path__ = [str(_shared_parent.parent)] +sys.modules['python_modules'] = _pm + +_pms = types.ModuleType('python_modules.shared') +_pms.__path__ = [str(_shared_parent)] +_pms.logger = _logger_module +sys.modules['python_modules.shared'] = _pms +sys.modules['python_modules.shared.logger'] = _logger_module + +# Register the real rate_limiter package +_rl = types.ModuleType('python_modules.shared.rate_limiter') +_rl.__path__ = [str(_rl_path)] +_rl.__file__ = str(_rl_path / "__init__.py") +sys.modules['python_modules.shared.rate_limiter'] = _rl + +# Load each rate_limiter submodule +for _mod_name in ('TokenBucket', 'DynamoDBMonitor', 'DistributedDynamoDBMonitorWorker', 'DistributedDynamoDBMonitorAggregator'): + _fqn = f'python_modules.shared.rate_limiter.{_mod_name}' + _spec = importlib.util.spec_from_file_location(_fqn, str(_rl_path / f'{_mod_name}.py')) + _mod = importlib.util.module_from_spec(_spec) + sys.modules[_fqn] = _mod + _spec.loader.exec_module(_mod) + setattr(_rl, _mod_name, _mod) diff --git a/tools/bulk_executor/tests/server/rate_limiter/test_distributed_aggregator.py b/tools/bulk_executor/tests/server/rate_limiter/test_distributed_aggregator.py new file mode 100644 index 00000000..c4618bed --- /dev/null +++ b/tools/bulk_executor/tests/server/rate_limiter/test_distributed_aggregator.py @@ -0,0 +1,232 @@ +"""Unit tests for DistributedDynamoDBMonitorAggregator — S3 aggregation, staleness cutoff, cleanup.""" +import json +import time +from io import BytesIO +from unittest.mock import Mock, MagicMock, patch, call +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from python_modules.shared.rate_limiter.DistributedDynamoDBMonitorAggregator import DistributedDynamoDBMonitorAggregator + + +@pytest.fixture +def mock_session(): + session = Mock() + s3 = Mock() + session.client = Mock(return_value=s3) + return session + + +@pytest.fixture +def s3_client(mock_session): + return mock_session.client.return_value + + +@pytest.fixture +def aggregator(mock_session, s3_client): + # Set up paginator to return empty by default + paginator = Mock() + paginator.paginate = Mock(return_value=[{"Contents": []}]) + s3_client.get_paginator = Mock(return_value=paginator) + + agg = DistributedDynamoDBMonitorAggregator( + session=mock_session, + bucket='test-bucket', + prefix='rate-limiter/job-123', + staleness_cutoff=15, + interval=60, # long so background thread doesn't fire + autostart=False, + ) + yield agg + agg.stop() + + +class TestInit: + def test_prefix_gets_trailing_slash(self, mock_session): + s3 = mock_session.client.return_value + paginator = Mock() + paginator.paginate = Mock(return_value=[]) + s3.get_paginator = Mock(return_value=paginator) + agg = DistributedDynamoDBMonitorAggregator( + session=mock_session, bucket='b', prefix='no-slash', + autostart=False, + ) + assert agg.prefix == 'no-slash/' + agg.stop() + + def test_preserves_trailing_slash(self, mock_session): + s3 = mock_session.client.return_value + paginator = Mock() + paginator.paginate = Mock(return_value=[]) + s3.get_paginator = Mock(return_value=paginator) + agg = DistributedDynamoDBMonitorAggregator( + session=mock_session, bucket='b', prefix='has/', + autostart=False, + ) + assert agg.prefix == 'has/' + agg.stop() + + +class TestAggregateOnce: + def test_writes_summary_with_zero_workers(self, aggregator, s3_client): + aggregator.aggregate_once() + s3_client.put_object.assert_called_once() + call_kwargs = s3_client.put_object.call_args[1] + summary = json.loads(call_kwargs['Body'].decode('utf-8')) + assert summary['active_workers'] == 0 + assert summary['aggregated_read_rate'] == 0.0 + assert summary['aggregated_write_rate'] == 0.0 + + def test_aggregates_multiple_workers(self, aggregator, s3_client): + now = time.time() + worker_data = [ + {"worker_id": "w1", "timestamp": now, "read_rate": 100.0, "write_rate": 50.0}, + {"worker_id": "w2", "timestamp": now, "read_rate": 200.0, "write_rate": 75.0}, + ] + + paginator = Mock() + paginator.paginate = Mock(return_value=[{ + "Contents": [ + {"Key": "rate-limiter/job-123/worker-w1.json"}, + {"Key": "rate-limiter/job-123/worker-w2.json"}, + ] + }]) + s3_client.get_paginator = Mock(return_value=paginator) + + def mock_get_object(Bucket, Key): + idx = 0 if 'w1' in Key else 1 + return {'Body': BytesIO(json.dumps(worker_data[idx]).encode('utf-8'))} + + s3_client.get_object = Mock(side_effect=mock_get_object) + + aggregator.aggregate_once() + + call_kwargs = s3_client.put_object.call_args[1] + summary = json.loads(call_kwargs['Body'].decode('utf-8')) + assert summary['active_workers'] == 2 + assert summary['aggregated_read_rate'] == pytest.approx(300.0) + assert summary['aggregated_write_rate'] == pytest.approx(125.0) + + def test_skips_stale_workers(self, aggregator, s3_client): + stale_ts = time.time() - 30 # 30s old, cutoff is 15 + fresh_ts = time.time() + + paginator = Mock() + paginator.paginate = Mock(return_value=[{ + "Contents": [ + {"Key": "rate-limiter/job-123/worker-stale.json"}, + {"Key": "rate-limiter/job-123/worker-fresh.json"}, + ] + }]) + s3_client.get_paginator = Mock(return_value=paginator) + + def mock_get_object(Bucket, Key): + if 'stale' in Key: + data = {"worker_id": "stale", "timestamp": stale_ts, "read_rate": 999.0, "write_rate": 999.0} + else: + data = {"worker_id": "fresh", "timestamp": fresh_ts, "read_rate": 50.0, "write_rate": 25.0} + return {'Body': BytesIO(json.dumps(data).encode('utf-8'))} + + s3_client.get_object = Mock(side_effect=mock_get_object) + + aggregator.aggregate_once() + + call_kwargs = s3_client.put_object.call_args[1] + summary = json.loads(call_kwargs['Body'].decode('utf-8')) + assert summary['active_workers'] == 1 + assert summary['aggregated_read_rate'] == pytest.approx(50.0) + + def test_skips_summary_file_in_listing(self, aggregator, s3_client): + paginator = Mock() + paginator.paginate = Mock(return_value=[{ + "Contents": [ + {"Key": "rate-limiter/job-123/summary.json"}, + {"Key": "rate-limiter/job-123/worker-w1.json"}, + ] + }]) + s3_client.get_paginator = Mock(return_value=paginator) + + fresh_ts = time.time() + s3_client.get_object = Mock(return_value={ + 'Body': BytesIO(json.dumps({ + "worker_id": "w1", "timestamp": fresh_ts, + "read_rate": 10.0, "write_rate": 5.0 + }).encode('utf-8')) + }) + + aggregator.aggregate_once() + + # get_object should only be called for worker file, not summary + assert s3_client.get_object.call_count == 1 + + def test_handles_invalid_timestamp(self, aggregator, s3_client): + paginator = Mock() + paginator.paginate = Mock(return_value=[{ + "Contents": [{"Key": "rate-limiter/job-123/worker-bad.json"}] + }]) + s3_client.get_paginator = Mock(return_value=paginator) + s3_client.get_object = Mock(return_value={ + 'Body': BytesIO(json.dumps({ + "worker_id": "bad", "timestamp": "not-a-number", + "read_rate": 100.0, "write_rate": 50.0 + }).encode('utf-8')) + }) + + aggregator.aggregate_once() + + call_kwargs = s3_client.put_object.call_args[1] + summary = json.loads(call_kwargs['Body'].decode('utf-8')) + assert summary['active_workers'] == 0 + + def test_handles_json_decode_error(self, aggregator, s3_client): + paginator = Mock() + paginator.paginate = Mock(return_value=[{ + "Contents": [{"Key": "rate-limiter/job-123/worker-corrupt.json"}] + }]) + s3_client.get_paginator = Mock(return_value=paginator) + s3_client.get_object = Mock(return_value={ + 'Body': BytesIO(b'not json at all') + }) + + aggregator.aggregate_once() + + call_kwargs = s3_client.put_object.call_args[1] + summary = json.loads(call_kwargs['Body'].decode('utf-8')) + assert summary['active_workers'] == 0 + + +class TestCleanup: + def test_cleanup_stops_and_deletes_summary(self, aggregator, s3_client): + aggregator.cleanup() + expected_key = f"{aggregator.prefix}{aggregator.output_key}" + s3_client.delete_object.assert_called_once_with( + Bucket='test-bucket', Key=expected_key + ) + + def test_cleanup_handles_delete_failure(self, aggregator, s3_client): + s3_client.delete_object = Mock(side_effect=Exception("denied")) + aggregator.cleanup() # should not raise + + +class TestStartStop: + def test_start_creates_thread(self, aggregator): + aggregator.start() + assert aggregator._thread is not None + assert aggregator._thread.is_alive() + aggregator.stop() + + def test_start_is_idempotent(self, aggregator): + aggregator.start() + first_thread = aggregator._thread + aggregator.start() + assert aggregator._thread is first_thread + aggregator.stop() + + def test_stop_joins_thread(self, aggregator): + aggregator.start() + aggregator.stop() + assert not aggregator._thread.is_alive() + + def test_stop_without_start_is_noop(self, aggregator): + aggregator.stop() # should not raise diff --git a/tools/bulk_executor/tests/server/rate_limiter/test_distributed_worker.py b/tools/bulk_executor/tests/server/rate_limiter/test_distributed_worker.py new file mode 100644 index 00000000..03256338 --- /dev/null +++ b/tools/bulk_executor/tests/server/rate_limiter/test_distributed_worker.py @@ -0,0 +1,254 @@ +"""Unit tests for DistributedDynamoDBMonitorWorker — S3 sync loop, scaling logic, staleness, cleanup.""" +import json +import threading +import time +from io import BytesIO +from unittest.mock import Mock, MagicMock, patch, call + +import pytest + +from python_modules.shared.rate_limiter.DistributedDynamoDBMonitorWorker import DistributedDynamoDBMonitorWorker + + +@pytest.fixture +def mock_session(): + session = Mock() + session.events = Mock() + session.events.register = Mock() + s3 = Mock() + session.client = Mock(return_value=s3) + return session + + +@pytest.fixture +def s3_client(mock_session): + return mock_session.client.return_value + + +@pytest.fixture +def worker(mock_session, s3_client): + s3_client.exceptions = Mock() + s3_client.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {}) + s3_client.get_object = Mock(side_effect=s3_client.exceptions.NoSuchKey("no summary")) + w = DistributedDynamoDBMonitorWorker( + session=mock_session, + bucket='test-bucket', + prefix='rate-limiter/job-123', + worker_max_read_rate=1500, + worker_max_write_rate=500, + sync_interval=60, # long interval to prevent auto-sync during tests + autostart=False, + ) + yield w + w.stop() + + +class TestInit: + def test_prefix_gets_trailing_slash(self, mock_session): + s3 = mock_session.client.return_value + s3.exceptions = Mock() + s3.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {}) + s3.get_object = Mock(side_effect=s3.exceptions.NoSuchKey("x")) + w = DistributedDynamoDBMonitorWorker( + session=mock_session, bucket='b', prefix='no-slash', + autostart=False, + ) + assert w.prefix == 'no-slash/' + w.stop() + + def test_prefix_preserved_if_has_slash(self, mock_session): + s3 = mock_session.client.return_value + s3.exceptions = Mock() + s3.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {}) + s3.get_object = Mock(side_effect=s3.exceptions.NoSuchKey("x")) + w = DistributedDynamoDBMonitorWorker( + session=mock_session, bucket='b', prefix='has-slash/', + autostart=False, + ) + assert w.prefix == 'has-slash/' + w.stop() + + def test_default_initial_rates(self, mock_session): + s3 = mock_session.client.return_value + s3.exceptions = Mock() + s3.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {}) + s3.get_object = Mock(side_effect=s3.exceptions.NoSuchKey("x")) + w = DistributedDynamoDBMonitorWorker( + session=mock_session, bucket='b', prefix='p/', + aggregate_max_read_rate=10000, + aggregate_max_write_rate=5000, + worker_max_read_rate=1500, + worker_max_write_rate=500, + autostart=False, + ) + # initial = min(worker_max, aggregate_max / 10) + assert w.monitor.max_read_rate == 1000.0 # min(1500, 10000/10) + assert w.monitor.max_write_rate == 500.0 # min(500, 5000/10) + w.stop() + + def test_custom_worker_id(self, mock_session): + s3 = mock_session.client.return_value + s3.exceptions = Mock() + s3.exceptions.NoSuchKey = type('NoSuchKey', (Exception,), {}) + s3.get_object = Mock(side_effect=s3.exceptions.NoSuchKey("x")) + w = DistributedDynamoDBMonitorWorker( + session=mock_session, bucket='b', prefix='p/', + worker_id='my-worker', autostart=False, + ) + assert w.worker_id == 'my-worker' + w.stop() + + +class TestSyncLoop: + def test_uploads_metrics_to_s3(self, worker, s3_client): + # Simulate one sync cycle manually + worker._last_metrics_snapshot = None + worker._sync_loop_once() + s3_client.put_object.assert_called_once() + call_kwargs = s3_client.put_object.call_args[1] + assert call_kwargs['Bucket'] == 'test-bucket' + assert 'worker-' in call_kwargs['Key'] + payload = json.loads(call_kwargs['Body'].decode('utf-8')) + assert 'worker_id' in payload + assert 'timestamp' in payload + assert 'read_rate' in payload + assert 'write_rate' in payload + + def test_computes_rate_from_delta(self, worker, s3_client): + # Set initial snapshot + worker._last_metrics_snapshot = (time.monotonic() - 5.0, 100.0, 50.0) + # Set current metrics + with worker.monitor.metrics_lock: + worker.monitor.metrics['read_capacity'] = 200.0 + worker.monitor.metrics['write_capacity'] = 100.0 + + worker._sync_loop_once() + payload = json.loads(s3_client.put_object.call_args[1]['Body'].decode('utf-8')) + # (200-100)/5 = 20, (100-50)/5 = 10 + assert payload['read_rate'] == pytest.approx(20.0, abs=1) + assert payload['write_rate'] == pytest.approx(10.0, abs=1) + + def test_first_sync_reports_zero_rate(self, worker, s3_client): + worker._last_metrics_snapshot = None + worker._sync_loop_once() + payload = json.loads(s3_client.put_object.call_args[1]['Body'].decode('utf-8')) + assert payload['read_rate'] == 0.0 + assert payload['write_rate'] == 0.0 + + +class TestScalingLogic: + def test_applies_scaling_from_summary(self, worker, s3_client): + summary = { + "aggregated_read_rate": 3000.0, + "aggregated_write_rate": 1000.0, + "active_workers": 3, + } + s3_client.get_object = Mock(return_value={ + 'Body': BytesIO(json.dumps(summary).encode('utf-8')) + }) + worker.monitor.max_read_rate = 1000.0 + worker.monitor.max_write_rate = 500.0 + worker._last_metrics_snapshot = (time.monotonic() - 5, 0, 0) + + worker._sync_loop_once() + + # aggregate_max_read=100000, agg_rate=3000 → scale=33.33 + # allowed = 33.33 * 1000 = 33333, capped at worker_max=1500 + # smoothed = 0.6*1000 + 0.4*1500 = 1200 + assert worker.monitor.max_read_rate == pytest.approx(1200.0, abs=5) + + def test_no_summary_file_keeps_rates(self, worker, s3_client): + s3_client.get_object = Mock( + side_effect=s3_client.exceptions.NoSuchKey("not found") + ) + original_read = worker.monitor.max_read_rate + original_write = worker.monitor.max_write_rate + worker._last_metrics_snapshot = (time.monotonic() - 5, 0, 0) + + worker._sync_loop_once() + + assert worker.monitor.max_read_rate == original_read + assert worker.monitor.max_write_rate == original_write + + +class TestCleanup: + def test_cleanup_stops_and_deletes_s3_file(self, worker, s3_client): + worker.cleanup() + expected_key = f"{worker.prefix}worker-{worker.worker_id}.json" + s3_client.delete_object.assert_called_once_with( + Bucket='test-bucket', Key=expected_key + ) + + def test_cleanup_handles_delete_failure(self, worker, s3_client, capsys): + s3_client.delete_object = Mock(side_effect=Exception("access denied")) + worker.cleanup() # should not raise + captured = capsys.readouterr() + assert "Warning" in captured.out or "failed" in captured.out + + +class TestStartStop: + def test_start_creates_thread(self, worker): + worker.start() + assert worker._sync_thread is not None + assert worker._sync_thread.is_alive() + worker.stop() + + def test_start_is_idempotent(self, worker): + worker.start() + first_thread = worker._sync_thread + worker.start() + assert worker._sync_thread is first_thread + worker.stop() + + def test_stop_joins_thread(self, worker): + worker.start() + worker.stop() + assert not worker._sync_thread.is_alive() + + +# Helper: expose single sync iteration for testing without threading +def _add_sync_loop_once(cls): + """Monkeypatch to run one iteration of the sync loop for deterministic testing.""" + def _sync_loop_once(self): + import json, time, random + upload_key = f"{self.prefix}worker-{self.worker_id}.json" + wall_ts = time.time() + mono_now = time.monotonic() + with self.monitor.metrics_lock: + total_read = float(self.monitor.metrics['read_capacity']) + total_write = float(self.monitor.metrics['write_capacity']) + if self._last_metrics_snapshot is None: + read_rate = 0.0 + write_rate = 0.0 + else: + last_t, last_r, last_w = self._last_metrics_snapshot + dt = max(mono_now - last_t, 1e-6) + read_rate = max(0.0, (total_read - last_r) / dt) + write_rate = max(0.0, (total_write - last_w) / dt) + self._last_metrics_snapshot = (mono_now, total_read, total_write) + payload = json.dumps({ + "worker_id": self.worker_id, + "timestamp": wall_ts, + "read_rate": read_rate, + "write_rate": write_rate, + }) + self.s3_client.put_object(Bucket=self.bucket, Key=upload_key, Body=payload.encode('utf-8')) + try: + resp = self.s3_client.get_object(Bucket=self.bucket, Key=f"{self.prefix}{self.summary_key}") + summary_data = json.loads(resp['Body'].read().decode('utf-8')) + current_agg_read_rate = summary_data.get("aggregated_read_rate", 0.0) + current_agg_write_rate = summary_data.get("aggregated_write_rate", 0.0) + read_scale = self.aggregate_max_read_rate / current_agg_read_rate if current_agg_read_rate > 0 else 1.0 + write_scale = self.aggregate_max_write_rate / current_agg_write_rate if current_agg_write_rate > 0 else 1.0 + worker_allowed_read_rate = read_scale * self.monitor.max_read_rate + worker_allowed_write_rate = write_scale * self.monitor.max_write_rate + new_read_target = min(worker_allowed_read_rate, self.worker_max_read_rate) + new_write_target = min(worker_allowed_write_rate, self.worker_max_write_rate) + smoothing_factor = 0.4 + self.monitor.max_read_rate = (1 - smoothing_factor) * self.monitor.max_read_rate + smoothing_factor * new_read_target + self.monitor.max_write_rate = (1 - smoothing_factor) * self.monitor.max_write_rate + smoothing_factor * new_write_target + except self.s3_client.exceptions.NoSuchKey: + pass + cls._sync_loop_once = _sync_loop_once + +_add_sync_loop_once(DistributedDynamoDBMonitorWorker) diff --git a/tools/bulk_executor/tests/server/rate_limiter/test_dynamodb_monitor.py b/tools/bulk_executor/tests/server/rate_limiter/test_dynamodb_monitor.py new file mode 100644 index 00000000..afc36edf --- /dev/null +++ b/tools/bulk_executor/tests/server/rate_limiter/test_dynamodb_monitor.py @@ -0,0 +1,206 @@ +"""Unit tests for DynamoDBMonitor — event hooks, capacity tracking, bucket deduction.""" +import threading +import time +from unittest.mock import Mock, MagicMock, patch, call + +import pytest + +from python_modules.shared.rate_limiter.DynamoDBMonitor import DynamoDBMonitor + + +@pytest.fixture +def mock_session(): + session = Mock() + session.events = Mock() + session.events.register = Mock() + return session + + +@pytest.fixture +def monitor(mock_session): + m = DynamoDBMonitor(mock_session, max_read_rate=1000, max_write_rate=500, enable_reporting=False) + yield m + m.stop() + + +class TestInit: + def test_registers_three_event_hooks(self, mock_session): + DynamoDBMonitor(mock_session, max_read_rate=100, max_write_rate=50, enable_reporting=False) + calls = mock_session.events.register.call_args_list + assert len(calls) == 3 + patterns = [c[0][0] for c in calls] + assert 'provide-client-params.dynamodb.*' in patterns + assert 'before-call.dynamodb.*' in patterns + assert 'after-call.dynamodb.*' in patterns + + def test_invalid_read_rate_raises(self, mock_session): + with pytest.raises(ValueError, match="Invalid rate limits"): + DynamoDBMonitor(mock_session, max_read_rate=0, max_write_rate=50, enable_reporting=False) + + def test_invalid_write_rate_raises(self, mock_session): + with pytest.raises(ValueError, match="Invalid rate limits"): + DynamoDBMonitor(mock_session, max_read_rate=100, max_write_rate=0, enable_reporting=False) + + def test_bucket_rates_match_config(self, mock_session): + m = DynamoDBMonitor(mock_session, max_read_rate=200, max_write_rate=100, enable_reporting=False) + assert m._read_bucket.rate == 200.0 + assert m._write_bucket.rate == 100.0 + m.stop() + + def test_bucket_capacity_is_rate_times_multiplier(self, mock_session): + m = DynamoDBMonitor(mock_session, max_read_rate=200, max_write_rate=100, enable_reporting=False) + assert m._read_bucket.capacity == 400.0 # 200 * 2 + assert m._write_bucket.capacity == 200.0 # 100 * 2 + m.stop() + + +class TestAddReturnConsumedCapacity: + def test_adds_param_when_missing(self, monitor): + params = {} + monitor._add_return_consumed_capacity(params) + assert params['ReturnConsumedCapacity'] == 'TOTAL' + + def test_preserves_existing_param(self, monitor): + params = {'ReturnConsumedCapacity': 'INDEXES'} + monitor._add_return_consumed_capacity(params) + assert params['ReturnConsumedCapacity'] == 'INDEXES' + + +class TestEnforceRateLimit: + def test_read_operations_use_read_bucket(self, monitor): + for op in ('GetItem', 'BatchGetItem', 'Query', 'Scan', 'TransactGetItems'): + model = Mock() + model.name = op + monitor._read_bucket.deduct(monitor._read_bucket.capacity * 3) + with patch.object(monitor._read_bucket, 'wait_until_positive') as mock_wait: + monitor._enforce_rate_limit(params={}, model=model) + mock_wait.assert_called_once() + + def test_write_operations_use_write_bucket(self, monitor): + for op in ('PutItem', 'UpdateItem', 'DeleteItem', 'BatchWriteItem', 'TransactWriteItems'): + model = Mock() + model.name = op + with patch.object(monitor._write_bucket, 'wait_until_positive') as mock_wait: + monitor._enforce_rate_limit(params={}, model=model) + mock_wait.assert_called_once() + + def test_other_operations_do_not_block(self, monitor): + model = Mock() + model.name = 'DescribeTable' + with patch.object(monitor._read_bucket, 'wait_until_positive') as mock_read: + with patch.object(monitor._write_bucket, 'wait_until_positive') as mock_write: + monitor._enforce_rate_limit(params={}, model=model) + mock_read.assert_not_called() + mock_write.assert_not_called() + + +class TestTrackConsumedCapacity: + def test_tracks_dict_consumed_capacity(self, monitor): + model = Mock() + model.name = 'PutItem' + parsed = { + 'ConsumedCapacity': { + 'CapacityUnits': 5.0, + 'WriteCapacityUnits': 5.0, + 'ReadCapacityUnits': 0.0, + } + } + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + assert monitor.metrics['write_capacity'] == 5.0 + assert monitor.metrics['read_capacity'] == 0.0 + + def test_tracks_list_consumed_capacity(self, monitor): + model = Mock() + model.name = 'BatchWriteItem' + parsed = { + 'ConsumedCapacity': [ + {'CapacityUnits': 3.0, 'WriteCapacityUnits': 3.0, 'ReadCapacityUnits': 0.0}, + {'CapacityUnits': 2.0, 'WriteCapacityUnits': 2.0, 'ReadCapacityUnits': 0.0}, + ] + } + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + assert monitor.metrics['write_capacity'] == 5.0 + assert monitor.metrics['calls'] == 2 + + def test_no_consumed_capacity_is_noop(self, monitor): + model = Mock() + model.name = 'PutItem' + parsed = {} + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + assert monitor.metrics['calls'] == 0 + + def test_ambiguous_capacity_for_read_op(self, monitor): + model = Mock() + model.name = 'Scan' + parsed = { + 'ConsumedCapacity': { + 'CapacityUnits': 10.0, + } + } + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + assert monitor.metrics['read_capacity'] == 10.0 + assert monitor.metrics['write_capacity'] == 0.0 + + +class TestBucketDeduction: + def test_read_capacity_deducted_from_read_bucket(self, monitor): + model = Mock() + model.name = 'Scan' + initial_snap = monitor._read_bucket.snapshot() + parsed = { + 'ConsumedCapacity': { + 'ReadCapacityUnits': 25.0, + 'WriteCapacityUnits': 0.0, + 'CapacityUnits': 25.0, + } + } + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + after_snap = monitor._read_bucket.snapshot() + assert after_snap['tokens'] < initial_snap['tokens'] + + def test_write_capacity_deducted_from_write_bucket(self, monitor): + model = Mock() + model.name = 'PutItem' + initial_snap = monitor._write_bucket.snapshot() + parsed = { + 'ConsumedCapacity': { + 'WriteCapacityUnits': 10.0, + 'ReadCapacityUnits': 0.0, + 'CapacityUnits': 10.0, + } + } + monitor._track_consumed_capacity(http_response=Mock(), parsed=parsed, model=model) + after_snap = monitor._write_bucket.snapshot() + assert after_snap['tokens'] < initial_snap['tokens'] + + +class TestRateSetters: + def test_set_max_read_rate(self, monitor): + monitor.max_read_rate = 2000 + assert monitor.max_read_rate == 2000.0 + assert monitor._read_bucket.rate == 2000.0 + + def test_set_max_write_rate(self, monitor): + monitor.max_write_rate = 800 + assert monitor.max_write_rate == 800.0 + assert monitor._write_bucket.rate == 800.0 + + def test_set_invalid_read_rate_raises(self, monitor): + with pytest.raises(ValueError, match="read rate must be >= 1"): + monitor.max_read_rate = 0 + + def test_set_invalid_write_rate_raises(self, monitor): + with pytest.raises(ValueError, match="write rate must be >= 1"): + monitor.max_write_rate = 0 + + +class TestReporting: + def test_reporting_thread_starts_when_enabled(self, mock_session): + m = DynamoDBMonitor(mock_session, max_read_rate=100, max_write_rate=50, enable_reporting=True) + assert m._report_thread.is_alive() + m.stop() + assert not m._report_thread.is_alive() + + def test_stop_is_idempotent(self, monitor): + monitor.stop() + monitor.stop() # should not raise diff --git a/tools/bulk_executor/tests/server/rate_limiter/test_token_bucket.py b/tools/bulk_executor/tests/server/rate_limiter/test_token_bucket.py new file mode 100644 index 00000000..9bff1ce3 --- /dev/null +++ b/tools/bulk_executor/tests/server/rate_limiter/test_token_bucket.py @@ -0,0 +1,218 @@ +"""Unit tests for TokenBucket — refill, deduct, wait_until_positive, reconfigure, negative balance, concurrency.""" +import threading +import time +from unittest.mock import patch + +import pytest + +from python_modules.shared.rate_limiter.TokenBucket import TokenBucket + + +class TestTokenBucketInit: + def test_defaults_rate_as_initial_and_capacity(self): + tb = TokenBucket(rate=100) + snap = tb.snapshot() + assert snap["rate"] == 100.0 + assert snap["capacity"] == 100.0 + assert snap["tokens"] == pytest.approx(100.0, abs=1) + + def test_custom_initial_and_capacity(self): + tb = TokenBucket(rate=50, initial=10, capacity=200) + snap = tb.snapshot() + assert snap["rate"] == 50.0 + assert snap["capacity"] == 200.0 + assert snap["tokens"] == pytest.approx(10.0, abs=1) + + def test_initial_clamped_to_capacity(self): + tb = TokenBucket(rate=10, initial=999, capacity=20) + snap = tb.snapshot() + assert snap["tokens"] == pytest.approx(20.0, abs=1) + + def test_zero_rate_raises(self): + with pytest.raises(ValueError, match="rate must be > 0"): + TokenBucket(rate=0) + + def test_negative_rate_raises(self): + with pytest.raises(ValueError, match="rate must be > 0"): + TokenBucket(rate=-5) + + +class TestRefill: + def test_refill_adds_tokens_over_time(self): + tb = TokenBucket(rate=1000, initial=0, capacity=2000) + time.sleep(0.05) + snap = tb.snapshot() + assert snap["tokens"] > 0 + assert snap["tokens"] <= 2000 + + def test_refill_capped_at_capacity(self): + tb = TokenBucket(rate=100000, initial=100000, capacity=100000) + time.sleep(0.01) + snap = tb.snapshot() + assert snap["tokens"] == pytest.approx(100000.0, abs=1) + + +class TestDeduct: + def test_deduct_reduces_tokens(self): + tb = TokenBucket(rate=100, initial=100, capacity=200) + tb.deduct(50) + snap = tb.snapshot() + assert snap["tokens"] < 100 + + def test_deduct_allows_negative_balance(self): + tb = TokenBucket(rate=10, initial=10, capacity=20) + tb.deduct(30) + snap = tb.snapshot() + assert snap["tokens"] < 0 + + def test_deduct_negative_amount_is_noop(self): + tb = TokenBucket(rate=100, initial=100, capacity=100) + tb.deduct(-50) + snap = tb.snapshot() + assert snap["tokens"] == pytest.approx(100.0, abs=2) + + def test_deduct_zero_is_noop(self): + tb = TokenBucket(rate=100, initial=50, capacity=100) + tb.deduct(0) + snap = tb.snapshot() + assert snap["tokens"] == pytest.approx(50.0, abs=2) + + +class TestWaitUntilPositive: + def test_returns_immediately_when_positive(self): + tb = TokenBucket(rate=100, initial=50, capacity=100) + start = time.monotonic() + tb.wait_until_positive() + elapsed = time.monotonic() - start + assert elapsed < 0.01 + + def test_blocks_when_negative_then_returns(self): + tb = TokenBucket(rate=10000, initial=0, capacity=20000) + tb.deduct(50) + start = time.monotonic() + tb.wait_until_positive() + elapsed = time.monotonic() - start + assert elapsed < 0.1 # high rate so should unblock fast + + def test_blocks_respects_refill_rate(self): + tb = TokenBucket(rate=100, initial=0, capacity=200) + tb.deduct(5) + start = time.monotonic() + tb.wait_until_positive() + elapsed = time.monotonic() - start + # need 5 tokens at 100/sec = 50ms minimum + assert elapsed >= 0.04 + + +class TestReconfigure: + def test_change_rate(self): + tb = TokenBucket(rate=100, initial=50, capacity=200) + tb.reconfigure(rate=500) + snap = tb.snapshot() + assert snap["rate"] == 500.0 + assert snap["capacity"] == 200.0 + + def test_change_capacity_clamps_tokens(self): + tb = TokenBucket(rate=100, initial=100, capacity=200) + tb.reconfigure(capacity=50) + snap = tb.snapshot() + assert snap["capacity"] == 50.0 + assert snap["tokens"] <= 50.0 + + def test_scale_tokens_on_capacity_change(self): + tb = TokenBucket(rate=100, initial=80, capacity=100) + tb.reconfigure(capacity=200, scale_tokens=True) + snap = tb.snapshot() + # 80/100 * 200 = 160 + assert snap["tokens"] == pytest.approx(160.0, abs=5) + + def test_invalid_rate_raises(self): + tb = TokenBucket(rate=100) + with pytest.raises(ValueError, match="rate must be > 0"): + tb.reconfigure(rate=0) + + def test_invalid_capacity_raises(self): + tb = TokenBucket(rate=100) + with pytest.raises(ValueError, match="capacity must be > 0"): + tb.reconfigure(capacity=-1) + + def test_reconfigure_wakes_waiters(self): + tb = TokenBucket(rate=1, initial=0, capacity=10) + tb.deduct(100) # deeply negative + + woke = threading.Event() + + def waiter(): + tb.wait_until_positive() + woke.set() + + t = threading.Thread(target=waiter, daemon=True) + t.start() + time.sleep(0.02) + assert not woke.is_set() + # Reconfigure to very high rate so waiter unblocks + tb.reconfigure(rate=100000) + t.join(timeout=1.0) + assert woke.is_set() + + +class TestConcurrency: + def test_concurrent_deducts_are_consistent(self): + tb = TokenBucket(rate=0.001, initial=1000, capacity=1000) + n_threads = 10 + deducts_per_thread = 100 + + def deductor(): + for _ in range(deducts_per_thread): + tb.deduct(1) + + threads = [threading.Thread(target=deductor) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + snap = tb.snapshot() + expected = 1000 - (n_threads * deducts_per_thread) + assert snap["tokens"] == pytest.approx(expected, abs=2) + + def test_concurrent_wait_and_deduct(self): + tb = TokenBucket(rate=50000, initial=100, capacity=100000) + results = [] + + def wait_then_record(): + tb.wait_until_positive() + results.append(True) + + def deductor(): + for _ in range(50): + tb.deduct(1) + time.sleep(0.001) + + threads = [] + for _ in range(5): + threads.append(threading.Thread(target=wait_then_record)) + threads.append(threading.Thread(target=deductor)) + + for t in threads: + t.start() + for t in threads: + t.join(timeout=5.0) + + assert len(results) == 5 + + +class TestNegativeBalance: + def test_negative_initial_clamped_to_neg_capacity(self): + tb = TokenBucket(rate=100, initial=-500, capacity=100) + snap = tb.snapshot() + assert snap["tokens"] == pytest.approx(-100.0, abs=1) + + def test_deep_negative_recovers_with_refill(self): + tb = TokenBucket(rate=10000, initial=0, capacity=20000) + tb.deduct(500) + snap_before = tb.snapshot() + assert snap_before["tokens"] < 0 + time.sleep(0.1) + snap_after = tb.snapshot() + assert snap_after["tokens"] > snap_before["tokens"] From 9f287fd0f85a29f9810eff50c51ebc88b71c67a3 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:46 -0700 Subject: [PATCH 13/22] Force built-in role refresh on version mismatch (#8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * feat(bootstrap): force built-in role refresh on version mismatch (bu-i84) Previously, _add_glue_job_role returned early when a role already existed, never updating its policies. New tool versions that add required permissions (e.g. quotas_policy) were silently ignored for users who had already bootstrapped, causing permission-missing errors after upgrades. Now built-in roles are tagged with a BulkDynamoDBVersion tag. On each bootstrap run, the stored tag is compared against __version__; on mismatch (or missing tag), all managed policies and inline policies are re-applied and the tag is updated. This makes policy state converge to what the current version expects without requiring users to manually delete and recreate roles. Custom roles (--XRole ) are exempt from version tracking and policy management since they are user-controlled. --- .../client/src/infrastructure/bootstrap.py | 44 ++++++- .../tests/client/test_bootstrap.py | 113 +++++++++++++++++- 2 files changed, 152 insertions(+), 5 deletions(-) diff --git a/tools/bulk_executor/client/src/infrastructure/bootstrap.py b/tools/bulk_executor/client/src/infrastructure/bootstrap.py index b729b9f3..4e74be43 100644 --- a/tools/bulk_executor/client/src/infrastructure/bootstrap.py +++ b/tools/bulk_executor/client/src/infrastructure/bootstrap.py @@ -37,6 +37,8 @@ THIRD_PARTY_PYTHON_MODULES, ) +ROLE_VERSION_TAG_KEY = 'BulkDynamoDBVersion' + class BootstrapInfrastructure: def __init__(self, env_configs): @@ -75,6 +77,10 @@ def _get_role_name(self, args): role_id = READ_WRITE_ROLE_ID if is_write_access else READ_ONLY_ROLE_ID return f"{GLUE_JOB_ROOT_ROLE_NAME}-{role_id}-{self.aws_region}" # region definition for separate region specific permissioning + def _is_builtin_role(self, args): + role_param = args.get('XRole', '') + return not role_param or role_param in READ_WRITE_ROLE_TYPES + def _add_glue_job_role(self, args): log.info("Adding Glue Job role...") self._prompt_for_role(args) @@ -120,7 +126,8 @@ def _add_glue_job_role(self, args): ] } - # Create the role + # Create the role or check if it needs a policy refresh + role_exists = False try: response = self.iam_client.create_role( RoleName=role_name, @@ -129,13 +136,23 @@ def _add_glue_job_role(self, args): log.info(f"Bulk Executor Glue Job Role created: {role_name}") log.debug(f'Role ARN: {response["Role"]["Arn"]}') - except self.iam_client.exceptions.EntityAlreadyExistsException as e: + except self.iam_client.exceptions.EntityAlreadyExistsException: + role_exists = True log.info(f"Found Bulk Executor Glue Job Role: {role_name}") - return # Roles exists. No additional actions needed. except Exception as e: log.error(f'Unexpected error: {e}') exit(1) + if role_exists: + if not self._is_builtin_role(args): + log.info(f"Custom role {role_name} exists. Skipping policy management.") + return + existing_version = self._get_role_version_tag(role_name) + if existing_version == VERSION: + log.info(f"Role {role_name} is up to date (version {VERSION}).") + return + log.info(f"Role {role_name} version mismatch (tagged={existing_version}, current={VERSION}). Refreshing policies...") + policy_arns = [ 'arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole', # S3 permissions etc are handled here ] @@ -173,6 +190,27 @@ def _add_glue_job_role(self, args): log.error(f'Unexpected error: {e}') exit(1) + self._tag_role_version(role_name) + + def _get_role_version_tag(self, role_name): + try: + response = self.iam_client.list_role_tags(RoleName=role_name) + for tag in response.get('Tags', []): + if tag['Key'] == ROLE_VERSION_TAG_KEY: + return tag['Value'] + return None + except Exception: + return None + + def _tag_role_version(self, role_name): + try: + self.iam_client.tag_role( + RoleName=role_name, + Tags=[{'Key': ROLE_VERSION_TAG_KEY, 'Value': VERSION}] + ) + except Exception as e: + log.warning(f'Could not tag role {role_name} with version: {e}') + def _is_existing_role(self, role_name): try: self.iam_client.get_role(RoleName=role_name) diff --git a/tools/bulk_executor/tests/client/test_bootstrap.py b/tools/bulk_executor/tests/client/test_bootstrap.py index 060dda2f..ddc42959 100644 --- a/tools/bulk_executor/tests/client/test_bootstrap.py +++ b/tools/bulk_executor/tests/client/test_bootstrap.py @@ -244,6 +244,61 @@ class NoSuchEntityException(Exception): assert exc.value.code == 1 +# -- _get_role_version_tag / _tag_role_version / _is_builtin_role -------- + +class TestRoleVersionTag: + def test_returns_version_when_tag_present(self, bootstrap): + bootstrap.iam_client.list_role_tags.return_value = { + 'Tags': [ + {'Key': 'BulkDynamoDBVersion', 'Value': '1.2.3'}, + {'Key': 'OtherTag', 'Value': 'x'}, + ] + } + assert bootstrap._get_role_version_tag('MyRole') == '1.2.3' + + def test_returns_none_when_tag_absent(self, bootstrap): + bootstrap.iam_client.list_role_tags.return_value = { + 'Tags': [{'Key': 'OtherTag', 'Value': 'x'}] + } + assert bootstrap._get_role_version_tag('MyRole') is None + + def test_returns_none_on_empty_tags(self, bootstrap): + bootstrap.iam_client.list_role_tags.return_value = {'Tags': []} + assert bootstrap._get_role_version_tag('MyRole') is None + + def test_returns_none_on_exception(self, bootstrap): + bootstrap.iam_client.list_role_tags.side_effect = RuntimeError('boom') + assert bootstrap._get_role_version_tag('MyRole') is None + + def test_tag_role_version_calls_tag_role(self, bootstrap): + from __version__ import __version__ as VERSION + bootstrap._tag_role_version('MyRole') + bootstrap.iam_client.tag_role.assert_called_once_with( + RoleName='MyRole', + Tags=[{'Key': 'BulkDynamoDBVersion', 'Value': VERSION}] + ) + + def test_tag_role_version_swallows_error(self, bootstrap): + bootstrap.iam_client.tag_role.side_effect = RuntimeError('no perms') + bootstrap._tag_role_version('MyRole') # should not raise + + +class TestIsBuiltinRole: + def test_no_role_is_builtin(self, bootstrap): + assert bootstrap._is_builtin_role({}) is True + + def test_read_only_is_builtin(self, bootstrap): + from infrastructure.constants import ROLE_TYPE_READ_ONLY + assert bootstrap._is_builtin_role({'XRole': ROLE_TYPE_READ_ONLY}) is True + + def test_read_write_is_builtin(self, bootstrap): + from infrastructure.constants import ROLE_TYPE_READ_WRITE + assert bootstrap._is_builtin_role({'XRole': ROLE_TYPE_READ_WRITE}) is True + + def test_custom_role_is_not_builtin(self, bootstrap): + assert bootstrap._is_builtin_role({'XRole': 'MyCustomRole'}) is False + + # -- _is_write_access_enabled ------------------------------------------- class TestIsWriteAccessEnabled: @@ -316,9 +371,11 @@ def test_creates_role_and_attaches_read_write_policies(self, bootstrap): assert 'arn:aws:iam::aws:policy/AmazonDynamoDBFullAccess' in attached assert 'arn:aws:iam::aws:policy/AmazonDynamoDBReadOnlyAccess' not in attached - def test_role_already_exists_returns_without_attaching_policies(self, bootstrap): + def test_role_already_exists_with_current_version_skips_policy_refresh(self, bootstrap): from infrastructure.constants import ROLE_TYPE_READ_ONLY + from __version__ import __version__ as VERSION bootstrap._prompt_for_role = MagicMock() + bootstrap._get_role_version_tag = MagicMock(return_value=VERSION) class EntityAlreadyExistsException(Exception): pass @@ -329,10 +386,62 @@ class EntityAlreadyExistsException(Exception): bootstrap._add_glue_job_role({'XRole': ROLE_TYPE_READ_ONLY}) - # Early return — no policy attachments bootstrap.iam_client.attach_role_policy.assert_not_called() bootstrap.iam_client.put_role_policy.assert_not_called() + def test_role_already_exists_with_old_version_refreshes_policies(self, bootstrap): + from infrastructure.constants import ROLE_TYPE_READ_ONLY + bootstrap._prompt_for_role = MagicMock() + bootstrap._get_role_version_tag = MagicMock(return_value='old-version') + + class EntityAlreadyExistsException(Exception): + pass + bootstrap.iam_client.exceptions.EntityAlreadyExistsException = ( + EntityAlreadyExistsException + ) + bootstrap.iam_client.create_role.side_effect = EntityAlreadyExistsException() + + bootstrap._add_glue_job_role({'XRole': ROLE_TYPE_READ_ONLY}) + + assert bootstrap.iam_client.attach_role_policy.call_count == 2 + assert bootstrap.iam_client.put_role_policy.call_count == 2 + bootstrap.iam_client.tag_role.assert_called_once() + + def test_role_already_exists_with_no_version_tag_refreshes_policies(self, bootstrap): + from infrastructure.constants import ROLE_TYPE_READ_ONLY + bootstrap._prompt_for_role = MagicMock() + bootstrap._get_role_version_tag = MagicMock(return_value=None) + + class EntityAlreadyExistsException(Exception): + pass + bootstrap.iam_client.exceptions.EntityAlreadyExistsException = ( + EntityAlreadyExistsException + ) + bootstrap.iam_client.create_role.side_effect = EntityAlreadyExistsException() + + bootstrap._add_glue_job_role({'XRole': ROLE_TYPE_READ_ONLY}) + + assert bootstrap.iam_client.attach_role_policy.call_count == 2 + assert bootstrap.iam_client.put_role_policy.call_count == 2 + bootstrap.iam_client.tag_role.assert_called_once() + + def test_custom_role_exists_skips_policy_management(self, bootstrap): + bootstrap._prompt_for_role = MagicMock() + bootstrap._is_existing_role = MagicMock(return_value=True) + + class EntityAlreadyExistsException(Exception): + pass + bootstrap.iam_client.exceptions.EntityAlreadyExistsException = ( + EntityAlreadyExistsException + ) + bootstrap.iam_client.create_role.side_effect = EntityAlreadyExistsException() + + bootstrap._add_glue_job_role({'XRole': 'MyCustomRole'}) + + bootstrap.iam_client.attach_role_policy.assert_not_called() + bootstrap.iam_client.put_role_policy.assert_not_called() + bootstrap.iam_client.tag_role.assert_not_called() + def test_unexpected_create_role_error_exits(self, bootstrap): from infrastructure.constants import ROLE_TYPE_READ_ONLY bootstrap._prompt_for_role = MagicMock() From 8cd70bc1ffc988d50bda3de888e0b0516113739f Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:50 -0700 Subject: [PATCH 14/22] Add --XIdleTimeout for Glue cost optimization (#9) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * feat: add --XIdleTimeout for Glue cost optimization (bu-i88) AWS Glue charges by DPU-minute. When workers finish processing but the job stays alive until the overall Timeout expires, the idle workers still incur cost. The IdleTimeout parameter (added to Glue's start_job_run API) causes the job to auto-terminate when all workers have been idle for the specified duration. Adds --XIdleTimeout (1-10080 minutes, default 5) as a configurable parameter following the existing XTimeout pattern. The 5-minute default balances cost savings against premature termination during inter-phase gaps in multi-step ETL jobs. Users processing small tables will see significant cost reduction since jobs no longer idle for the remaining 55 minutes of the default 60-minute timeout. The parameter is passed to start_job_run only (not create_job/update_job) because IdleTimeout is a run-level setting in the Glue API. --- .../client/src/infrastructure/constants.py | 1 + tools/bulk_executor/client/src/runner.py | 1 + tools/bulk_executor/client/src/utils/__init__.py | 12 ++++++++++++ tools/bulk_executor/tests/client/test_runner.py | 12 ++++++++++++ .../tests/client/utils/test_utils_init.py | 15 +++++++++++++++ 5 files changed, 41 insertions(+) diff --git a/tools/bulk_executor/client/src/infrastructure/constants.py b/tools/bulk_executor/client/src/infrastructure/constants.py index 0cee3093..4e9f7ad8 100644 --- a/tools/bulk_executor/client/src/infrastructure/constants.py +++ b/tools/bulk_executor/client/src/infrastructure/constants.py @@ -38,6 +38,7 @@ class GlueJobDefaults(Enum): MaxConcurrentRuns=20 Retries=0 Timeout=60 + IdleTimeout=5 NumberOfWorkers=220 WorkerType='G.1X' diff --git a/tools/bulk_executor/client/src/runner.py b/tools/bulk_executor/client/src/runner.py index 7f060a48..b969d5ae 100644 --- a/tools/bulk_executor/client/src/runner.py +++ b/tools/bulk_executor/client/src/runner.py @@ -374,6 +374,7 @@ def _start_glue_job(self, glue_job_arguments, args): NumberOfWorkers=args.get('XNumberOfWorkers', GlueJobDefaults.NumberOfWorkers.value), Timeout=args.get('XTimeout', GlueJobDefaults.Timeout.value), WorkerType=args.get('XWorkerType', GlueJobDefaults.WorkerType.value), + IdleTimeout=args.get('XIdleTimeout', GlueJobDefaults.IdleTimeout.value), ) return response['JobRunId'] except Exception as e: diff --git a/tools/bulk_executor/client/src/utils/__init__.py b/tools/bulk_executor/client/src/utils/__init__.py index c03e32df..28f1f636 100644 --- a/tools/bulk_executor/client/src/utils/__init__.py +++ b/tools/bulk_executor/client/src/utils/__init__.py @@ -81,6 +81,17 @@ def validate_timeout(x): raise argparse.ArgumentTypeError(f"Timeout must be an integer") +def validate_idle_timeout(x): + try: + value = int(x) + if 1 <= value <= 10080: + return value + else: + raise argparse.ArgumentTypeError(f"IdleTimeout must be between 1 and 10080 minutes (7 days)") + except ValueError: + raise argparse.ArgumentTypeError(f"IdleTimeout must be an integer") + + # The defaults stated below should perhaps be dynamic from constants.py def parse_action(): parser = argparse.ArgumentParser( @@ -98,6 +109,7 @@ def glue_job_arguments(): parser.add_argument("--XExecutionClass", type=str, default=argparse.SUPPRESS, help="Set to STANDARD (default) or FLEX (lower DPU cost by using spare capacity, may take longer).", choices=SUPPORTED_EXECUTION_CLASSES) parser.add_argument("--XTimeout", type=validate_timeout, default=argparse.SUPPRESS, help="The Glue Job timeout (in minutes). Must be between 1 and 10080 minutes (7 days is the max allowed timeout). Default is 60 minutes.") + parser.add_argument("--XIdleTimeout", type=validate_idle_timeout, default=argparse.SUPPRESS, help="The Glue Job idle timeout (in minutes). Workers stop early when idle for this duration, reducing cost. Must be between 1 and 10080 minutes. Default is 5 minutes.") parser.add_argument("--XNumberOfWorkers", type=int, default=argparse.SUPPRESS, help="The number of Glue workers (default 220).") parser.add_argument("--XWorkerType", type=str, default=argparse.SUPPRESS, help="The Glue worker type. ex. G.1X", choices=SUPPORTED_WORKER_TYPES) parser.add_argument("--XWaitForDPU", action='store_true', default=argparse.SUPPRESS, help="Causes execution to wait 40 seconds at the end of execution for DPU metrics to be available.") diff --git a/tools/bulk_executor/tests/client/test_runner.py b/tools/bulk_executor/tests/client/test_runner.py index d4707fa9..a7050c51 100644 --- a/tools/bulk_executor/tests/client/test_runner.py +++ b/tools/bulk_executor/tests/client/test_runner.py @@ -942,6 +942,18 @@ def test_client_error_with_unknown_code_exits(self, bulk_runner): with pytest.raises(SystemExit): bulk_runner._start_glue_job({}, {}) + def test_uses_default_idle_timeout_when_missing(self, bulk_runner): + bulk_runner.glue_client.start_job_run.return_value = {'JobRunId': 'x'} + bulk_runner._start_glue_job({}, {}) + kwargs = bulk_runner.glue_client.start_job_run.call_args.kwargs + assert kwargs['IdleTimeout'] == runner_module.GlueJobDefaults.IdleTimeout.value + + def test_overrides_idle_timeout_with_provided_arg(self, bulk_runner): + bulk_runner.glue_client.start_job_run.return_value = {'JobRunId': 'x'} + bulk_runner._start_glue_job({}, {'XIdleTimeout': 10}) + kwargs = bulk_runner.glue_client.start_job_run.call_args.kwargs + assert kwargs['IdleTimeout'] == 10 + # --- _stop_glue_job --------------------------------------------------------- diff --git a/tools/bulk_executor/tests/client/utils/test_utils_init.py b/tools/bulk_executor/tests/client/utils/test_utils_init.py index 59429a08..e5625804 100644 --- a/tools/bulk_executor/tests/client/utils/test_utils_init.py +++ b/tools/bulk_executor/tests/client/utils/test_utils_init.py @@ -243,6 +243,21 @@ def test_timeout_validation_rejects_out_of_range(self, capsys): with pytest.raises(SystemExit): parser.parse_known_args(['--XTimeout', '99999']) + def test_parses_idle_timeout_via_validator(self): + parser = glue_job_arguments() + ns, _ = parser.parse_known_args(['--XIdleTimeout', '10']) + assert ns.XIdleTimeout == 10 + + def test_idle_timeout_validation_rejects_out_of_range(self, capsys): + parser = glue_job_arguments() + with pytest.raises(SystemExit): + parser.parse_known_args(['--XIdleTimeout', '99999']) + + def test_idle_timeout_validation_rejects_zero(self, capsys): + parser = glue_job_arguments() + with pytest.raises(SystemExit): + parser.parse_known_args(['--XIdleTimeout', '0']) + def test_parses_number_of_workers(self): parser = glue_job_arguments() ns, _ = parser.parse_known_args(['--XNumberOfWorkers', '50']) From 30d117446742769b0e4eb202e54b4fc0a012123f Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:53 -0700 Subject: [PATCH 15/22] Add rate validation warnings for XMaxReadRate/XMaxWriteRate (#10) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * Add rate validation warnings for XMaxReadRate/XMaxWriteRate Warn users when their configured read/write rate is: - Too high: exceeds table provisioned capacity or on-demand limit - Too low: below minimum recommended or would make job unreasonably slow Includes suggested rate ranges based on actual table capacity. Fixes #89 --- .../src/python_modules/shared/table_info.py | 57 +++++++ .../tests/server/test_table_info.py | 148 ++++++++++++++++++ 2 files changed, 205 insertions(+) diff --git a/tools/bulk_executor/server/src/python_modules/shared/table_info.py b/tools/bulk_executor/server/src/python_modules/shared/table_info.py index 6c6d31aa..3d2cfe4b 100644 --- a/tools/bulk_executor/server/src/python_modules/shared/table_info.py +++ b/tools/bulk_executor/server/src/python_modules/shared/table_info.py @@ -11,6 +11,45 @@ MIN_RECOMMENDED_READ_RATE = 100 MIN_RECOMMENDED_WRITE_RATE = 100 +# Maximum multiplier over table capacity before we warn the user. +# For provisioned tables, exceeding 100% will cause throttling. +# For on-demand tables, exceeding the configured limit causes throttling. +MAX_RATE_CAPACITY_MULTIPLIER = 1.0 + + +def _warn_rate_vs_capacity(table_name, direction, user_rate, table_capacity): + """Emit a warning if user-specified rate is too high or too low relative to table capacity. + + Returns suggested (low, high) range tuple for testing. + """ + user_rate = int(user_rate) + table_capacity = int(table_capacity) + + min_recommended = MIN_RECOMMENDED_READ_RATE if direction == "read" else MIN_RECOMMENDED_WRITE_RATE + suggested_low = max(min_recommended, table_capacity // 10) + suggested_high = table_capacity + + if user_rate > table_capacity: + log.warn( + f"[{table_name}] WARNING: {direction.capitalize()} rate {user_rate:,} exceeds table capacity " + f"of {table_capacity:,}. This will cause throttling. " + f"Suggested range: {suggested_low:,}–{suggested_high:,}." + ) + elif user_rate < min_recommended: + log.warn( + f"[{table_name}] WARNING: {direction.capitalize()} rate {user_rate:,} is very low and will " + f"result in an extremely slow job. " + f"Suggested range: {suggested_low:,}–{suggested_high:,}." + ) + elif user_rate < suggested_low: + log.warn( + f"[{table_name}] WARNING: {direction.capitalize()} rate {user_rate:,} is low relative to table " + f"capacity of {table_capacity:,}. The job may take unreasonably long. " + f"Suggested range: {suggested_low:,}–{suggested_high:,}." + ) + + return (suggested_low, suggested_high) + def get_quota_value(quota_name, region_name): """ Get the value of a specific DynamoDB quota from Service Quotas API. @@ -327,6 +366,15 @@ def get_dynamodb_throughput_configs(args, table_name, modes=None, format="connec if "read" in modes: if read_rate: log.info(f"[{table_name}] Max read rate set to specified limit: {read_rate}") + # Warn if user-specified rate is unreasonable relative to table capacity + effective_read_capacity = None + if is_on_demand_table: + on_demand_throughput = table_desc.get('OnDemandThroughput', {}) + effective_read_capacity = on_demand_throughput.get('MaxReadRequestUnits') or DEFAULT_ON_DEMAND_CAPACITY + else: + effective_read_capacity = table_desc.get('ProvisionedThroughput', {}).get('ReadCapacityUnits') + if effective_read_capacity: + _warn_rate_vs_capacity(table_name, "read", read_rate, effective_read_capacity) elif is_on_demand_table: # Check for table-specific limit on_demand_throughput = table_desc.get('OnDemandThroughput', {}) @@ -360,6 +408,15 @@ def get_dynamodb_throughput_configs(args, table_name, modes=None, format="connec if "write" in modes: if write_rate: log.info(f"[{table_name}] Max write rate set to specified limit: {write_rate}") + # Warn if user-specified rate is unreasonable relative to table capacity + effective_write_capacity = None + if is_on_demand_table: + on_demand_throughput = table_desc.get('OnDemandThroughput', {}) + effective_write_capacity = on_demand_throughput.get('MaxWriteRequestUnits') or DEFAULT_ON_DEMAND_CAPACITY + else: + effective_write_capacity = table_desc.get('ProvisionedThroughput', {}).get('WriteCapacityUnits') + if effective_write_capacity: + _warn_rate_vs_capacity(table_name, "write", write_rate, effective_write_capacity) elif is_on_demand_table: # Check for table-specific limit on_demand_throughput = table_desc.get('OnDemandThroughput', {}) diff --git a/tools/bulk_executor/tests/server/test_table_info.py b/tools/bulk_executor/tests/server/test_table_info.py index ae2ec32c..92c727cb 100644 --- a/tools/bulk_executor/tests/server/test_table_info.py +++ b/tools/bulk_executor/tests/server/test_table_info.py @@ -1374,3 +1374,151 @@ def test_read_rate_zero_emits_only_percent_key(self, boto3_mock): # but the read percent key is always set. assert 'dynamodb.throughput.read' not in opts assert opts['dynamodb.throughput.read.percent'] == '1.0' + + +# --- _warn_rate_vs_capacity --------------------------------------------------- + + +class TestWarnRateVsCapacity: + """Unit tests for the _warn_rate_vs_capacity helper.""" + + def test_returns_suggested_range(self): + low, high = table_info._warn_rate_vs_capacity('t', 'read', 500, 1000) + assert low == 100 # max(MIN_RECOMMENDED=100, 1000//10=100) + assert high == 1000 + + def test_suggested_low_is_at_least_min_recommended(self): + low, high = table_info._warn_rate_vs_capacity('t', 'write', 50, 500) + # 500 // 10 = 50, but MIN_RECOMMENDED_WRITE_RATE=100 wins + assert low == 100 + assert high == 500 + + def test_suggested_low_scales_with_large_capacity(self): + low, high = table_info._warn_rate_vs_capacity('t', 'read', 5000, 40000) + # 40000 // 10 = 4000 > MIN=100 + assert low == 4000 + assert high == 40000 + + def test_too_high_emits_warning(self, caplog): + import logging + with caplog.at_level(logging.WARNING): + table_info._warn_rate_vs_capacity('t', 'read', 2000, 1000) + assert 'exceeds table capacity' in caplog.text + assert '2,000' in caplog.text + assert '1,000' in caplog.text + + def test_too_low_below_min_emits_warning(self, caplog): + import logging + with caplog.at_level(logging.WARNING): + table_info._warn_rate_vs_capacity('t', 'write', 5, 1000) + assert 'very low' in caplog.text + + def test_low_relative_to_capacity_emits_warning(self, caplog): + import logging + with caplog.at_level(logging.WARNING): + # capacity=40000, rate=200 is above MIN_RECOMMENDED but below 40000//10=4000 + table_info._warn_rate_vs_capacity('t', 'read', 200, 40000) + assert 'low relative to table capacity' in caplog.text + + def test_rate_within_range_no_warning(self, caplog): + import logging + with caplog.at_level(logging.WARNING): + table_info._warn_rate_vs_capacity('t', 'read', 500, 1000) + assert 'WARNING' not in caplog.text + assert 'exceeds' not in caplog.text + assert 'low' not in caplog.text + + +# --- Rate validation integration in get_dynamodb_throughput_configs ----------- + + +class TestRateValidationIntegration: + """Verify _warn_rate_vs_capacity is called when users specify explicit rates.""" + + def test_read_rate_too_high_provisioned_warns( + self, boto3_mock, provisioned_table, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = provisioned_table + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxReadRate': '5000'}, table_name='t', modes=['read'] + ) + # provisioned RCU=1000, user wants 5000 → too high + assert 'exceeds table capacity' in caplog.text + + def test_write_rate_too_high_provisioned_warns( + self, boto3_mock, provisioned_table, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = provisioned_table + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxWriteRate': '2000'}, table_name='t', modes=['write'] + ) + # provisioned WCU=500, user wants 2000 → too high + assert 'exceeds table capacity' in caplog.text + + def test_read_rate_too_high_ondemand_warns( + self, boto3_mock, ondemand_table_with_limits, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = ( + ondemand_table_with_limits + ) + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxReadRate': '50000'}, table_name='t', modes=['read'] + ) + # on-demand table read limit=25000, user wants 50000 → too high + assert 'exceeds table capacity' in caplog.text + + def test_write_rate_too_low_provisioned_warns( + self, boto3_mock, provisioned_table, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = provisioned_table + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxWriteRate': '10'}, table_name='t', modes=['write'] + ) + # provisioned WCU=500, user wants 10 → very low (below MIN_RECOMMENDED) + assert 'very low' in caplog.text + + def test_rate_within_capacity_no_capacity_warning( + self, boto3_mock, provisioned_table, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = provisioned_table + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxReadRate': '800'}, table_name='t', modes=['read'] + ) + # provisioned RCU=1000, user wants 800 → within range, no warning + assert 'exceeds' not in caplog.text + assert 'low relative' not in caplog.text + assert 'very low' not in caplog.text + + def test_no_capacity_warning_when_describe_table_fails( + self, boto3_mock, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.side_effect = RuntimeError('nope') + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxReadRate': '99999'}, table_name='t', modes=['read'] + ) + # describe_table failed → no capacity info → no capacity warning + assert 'exceeds table capacity' not in caplog.text + + def test_ondemand_no_table_limit_uses_default_40k( + self, boto3_mock, ondemand_table, caplog + ): + import logging + boto3_mock.dynamodb_client.describe_table.return_value = ondemand_table + with caplog.at_level(logging.WARNING): + table_info.get_dynamodb_throughput_configs( + args={'XMaxWriteRate': '80000'}, table_name='t', modes=['write'] + ) + # on-demand no table limit → default capacity=40000, user wants 80000 → too high + assert 'exceeds table capacity' in caplog.text From eb919d547d2b58872d79a6f6a1f85b7418ce65c3 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:00:57 -0700 Subject: [PATCH 16/22] Install faker only for fill verb to speed worker startup (#11) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * perf: install faker only for fill verb to speed worker startup (bu-i91) Worker logs showed 8-15 seconds spent installing faker via pip on every Glue job run, even for verbs (copy, find, delete, etc.) that never use it. The faker dependency was baked into the Glue job's DefaultArguments as --additional-python-modules, causing unconditional installation at worker spin-up. Move faker out of the global _THIRD_PARTY_PYTHON_MODULES list (which populates DefaultArguments at bootstrap time) and into a new VERB_PYTHON_MODULES dict keyed by verb name. The runner now injects --additional-python-modules into per-run Arguments only when the action matches a verb with extra dependencies (currently only 'fill'). This eliminates the pip install overhead for all other verbs while preserving fill's access to faker. Validated: 479 existing client+fill tests pass; manual assertions confirm faker is injected per-run for fill and omitted for copy/find. --- .../client/src/infrastructure/bootstrap.py | 4 +++- .../client/src/infrastructure/constants.py | 14 ++++++++++---- tools/bulk_executor/client/src/runner.py | 7 ++++++- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/tools/bulk_executor/client/src/infrastructure/bootstrap.py b/tools/bulk_executor/client/src/infrastructure/bootstrap.py index 4e74be43..ec60d72a 100644 --- a/tools/bulk_executor/client/src/infrastructure/bootstrap.py +++ b/tools/bulk_executor/client/src/infrastructure/bootstrap.py @@ -290,10 +290,12 @@ def _create_or_update_glue_job(self, args, is_create_allowed=True): '--s3-bucket-name': glue_job_bucket, '--s3-script-location': s3_script_location, '--extra-py-files': s3_python_module_location, - '--additional-python-modules': THIRD_PARTY_PYTHON_MODULES, '--bulk-dynamodb-version': VERSION }) + if THIRD_PARTY_PYTHON_MODULES: + default_arguments['--additional-python-modules'] = THIRD_PARTY_PYTHON_MODULES + # Custom log4j2 properties override continuous CloudWatch logging # Reference: https://repost.aws/knowledge-center/glue-reduce-cloudwatch-logs # "If you apply a custom log4j.properties or log4j2.properties config file, then AWS Glue turns off continuous logging" diff --git a/tools/bulk_executor/client/src/infrastructure/constants.py b/tools/bulk_executor/client/src/infrastructure/constants.py index 4e9f7ad8..1bdf1ed2 100644 --- a/tools/bulk_executor/client/src/infrastructure/constants.py +++ b/tools/bulk_executor/client/src/infrastructure/constants.py @@ -42,10 +42,16 @@ class GlueJobDefaults(Enum): NumberOfWorkers=220 WorkerType='G.1X' -# Third Party Dependencies as an alpha-numeric list -_THIRD_PARTY_PYTHON_MODULES = [ - 'faker' -] +# Third Party Dependencies installed on every Glue run (alpha-numeric list). +# Verb-specific deps go in VERB_PYTHON_MODULES below. +_THIRD_PARTY_PYTHON_MODULES = [] # Convert to AWS Glue Readable Format THIRD_PARTY_PYTHON_MODULES = ','.join(map(str, _THIRD_PARTY_PYTHON_MODULES)) + +# Dependencies needed only for specific verbs (keyed by verb name). +# These are added per-run via start_job_run Arguments to avoid +# installing them for every job invocation. +VERB_PYTHON_MODULES = { + 'fill': ['faker'], +} diff --git a/tools/bulk_executor/client/src/runner.py b/tools/bulk_executor/client/src/runner.py index b969d5ae..f02be645 100644 --- a/tools/bulk_executor/client/src/runner.py +++ b/tools/bulk_executor/client/src/runner.py @@ -14,7 +14,7 @@ # project files from clients import Clients from infrastructure import GLUE_JOB_NAME, GlueJobDefaults -from infrastructure.constants import GLUE_LOG_GROUP_ERROR, GLUE_LOG_GROUP_OUTPUT +from infrastructure.constants import GLUE_LOG_GROUP_ERROR, GLUE_LOG_GROUP_OUTPUT, VERB_PYTHON_MODULES from infrastructure.verifier import assert_version_parity, is_existing_glue_job from reassembler import GlueLogReassembler from utils.graceful_interrupt_handler import GracefulInterruptHandler @@ -334,6 +334,11 @@ def _get_glue_job_arguments(self, args, script_args): key = "XAction" arguments[f"--{key}"] = value + action = arguments.get('--XAction', '') + verb_modules = VERB_PYTHON_MODULES.get(action, []) + if verb_modules: + arguments['--additional-python-modules'] = ','.join(verb_modules) + log.debug(f"All Glue Job args: {arguments}") return arguments From c2044ff3aa8e92658ba8d9b3b939e7ba722857cc Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Sun, 28 Jun 2026 18:01:01 -0700 Subject: [PATCH 17/22] Add --XExistingBucket parameter to bootstrap (#13) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover JSON-array and nested-braces cases for _jsonify_message Addresses review feedback on PR #171 (coverage was thin — one happy-path test). Adds the second regex alternation branch ([...]) and a nested-object case. Full suite: 1275 passed, 99.7% line / 96.8% branch coverage. Co-authored-by: dark-factory-agent * feat: add --XExistingBucket parameter to bootstrap (GH#130) Allow users who cannot grant bucket-creation permissions to pass an existing S3 bucket to bootstrap. The bucket is validated for existence and accessibility before proceeding. Bucket creation is skipped but policy application and file uploads still occur. --- .../client/src/infrastructure/bootstrap.py | 53 ++++++--- .../client/src/python_modules/bootstrap.py | 1 + .../tests/client/test_bootstrap.py | 109 ++++++++++++++++++ 3 files changed, 145 insertions(+), 18 deletions(-) diff --git a/tools/bulk_executor/client/src/infrastructure/bootstrap.py b/tools/bulk_executor/client/src/infrastructure/bootstrap.py index ec60d72a..e6bd58ee 100644 --- a/tools/bulk_executor/client/src/infrastructure/bootstrap.py +++ b/tools/bulk_executor/client/src/infrastructure/bootstrap.py @@ -277,7 +277,7 @@ def _create_or_update_glue_job(self, args, is_create_allowed=True): # concerns and are not needed at job runtime, so they are excluded # from DefaultArguments. See issue #85. for key, value in args.items(): - if key.startswith('X') and key not in ('XRole', 'XRegion', 'XAccount'): + if key.startswith('X') and key not in ('XRole', 'XRegion', 'XAccount', 'XExistingBucket'): default_arguments[f'--{key}'] = str(value) default_arguments.update({ # Update last intentional. @@ -369,23 +369,24 @@ def _bucket_exists(self, s3_client, bucket_name): def _upload_job_root_to_s3(self): glue_job_bucket = self._get_glue_job_bucket_name() - # Check if the bucket exists - if not self._bucket_exists(self.s3_client, glue_job_bucket): - try: - # Create the bucket - bucket_config = {} - if self.aws_region != 'us-east-1': # Default is us-east-1 so LocationConstraint fails if configured for this region. - bucket_config['CreateBucketConfiguration'] = {'LocationConstraint': self.aws_region} - self.s3_client.create_bucket( - Bucket=glue_job_bucket, - **bucket_config - ) - log.info(f"Bucket '{glue_job_bucket}' created successfully!") - except Exception as e: - log.error(f"Error creating bucket '{glue_job_bucket}': {e}") - exit(1) - else: - log.info(f"Bucket '{glue_job_bucket}' already exists.") + if not getattr(self, '_existing_bucket', None): + # Check if the bucket exists + if not self._bucket_exists(self.s3_client, glue_job_bucket): + try: + # Create the bucket + bucket_config = {} + if self.aws_region != 'us-east-1': # Default is us-east-1 so LocationConstraint fails if configured for this region. + bucket_config['CreateBucketConfiguration'] = {'LocationConstraint': self.aws_region} + self.s3_client.create_bucket( + Bucket=glue_job_bucket, + **bucket_config + ) + log.info(f"Bucket '{glue_job_bucket}' created successfully!") + except Exception as e: + log.error(f"Error creating bucket '{glue_job_bucket}': {e}") + exit(1) + else: + log.info(f"Bucket '{glue_job_bucket}' already exists.") # Apply the secure transport policy try: @@ -423,7 +424,16 @@ def _upload_job_root_to_s3(self): self.s3_client.upload_file(f"./{GLUE_JOB_SERVER_ROOT_PATH}", glue_job_bucket, GLUE_JOB_SERVER_ROOT_PATH) log.info(f"Glue script '{GLUE_JOB_SERVER_ROOT_PATH}' uploaded into S3 successfully.") + def _validate_existing_bucket(self, bucket_name): + if not self._bucket_exists(self.s3_client, bucket_name): + log.error(f"The specified bucket '{bucket_name}' does not exist or is not accessible.") + exit(1) + log.info(f"Using existing S3 bucket: {bucket_name}") + def _get_glue_job_bucket_name(self): + if getattr(self, '_existing_bucket', None): + return self._existing_bucket + # Return the existing persisted S3 Bucket name job_details = self._get_glue_job_details() if job_details: @@ -588,6 +598,13 @@ def _create_glue_log_groups(self): exit(1) def bootstrap(self, args): + existing_bucket = args.get('XExistingBucket') + if existing_bucket: + self._validate_existing_bucket(existing_bucket) + self._existing_bucket = existing_bucket + else: + self._existing_bucket = None + self._add_glue_job_role(args) self._create_glue_log_groups() self._ensure_dynamodb_glue_connection() diff --git a/tools/bulk_executor/client/src/python_modules/bootstrap.py b/tools/bulk_executor/client/src/python_modules/bootstrap.py index a43c0f9c..cb9b09d6 100644 --- a/tools/bulk_executor/client/src/python_modules/bootstrap.py +++ b/tools/bulk_executor/client/src/python_modules/bootstrap.py @@ -40,6 +40,7 @@ def run(env_configs): parser.add_argument("--XRole", type=validate_role, default=argparse.SUPPRESS, help="The AWS Role to use for executing the action. Specify a custom role name, or use the special keywords READ-ONLY or READ-WRITE to generate a managed role.") parser.add_argument("--XMaxConcurrentRuns", type=int, default=argparse.SUPPRESS, help="The maximum number of concurrent runs to allow for the single Glue Job (default is 20).") parser.add_argument("--XRetries", type=int, default=argparse.SUPPRESS, help="The max number of Glue Job retries. Defaults to zero to fail a misconfigured Glue Job quickly.") + parser.add_argument("--XExistingBucket", type=str, default=argparse.SUPPRESS, help="Use an existing S3 bucket instead of creating a new one. The bucket must exist and be accessible. It should follow the Glue naming prefix convention (aws-glue-*).") #result = {k: v for k, v in vars(args).items() if v is not None} result = parser.parse_args().__dict__ diff --git a/tools/bulk_executor/tests/client/test_bootstrap.py b/tools/bulk_executor/tests/client/test_bootstrap.py index ddc42959..cb066c64 100644 --- a/tools/bulk_executor/tests/client/test_bootstrap.py +++ b/tools/bulk_executor/tests/client/test_bootstrap.py @@ -970,3 +970,112 @@ def test_init_wires_clients_and_env(self): assert instance.s3_client is clients.s3_client assert instance.glue_client is clients.glue_client assert instance.logs_client is clients.logs_client + + +# -- XExistingBucket feature (GH#130) ------------------------------------- + +class TestExistingBucket: + """Coverage for the --XExistingBucket parameter that allows users to + skip bucket creation and use a pre-existing S3 bucket.""" + + def test_existing_bucket_is_used_by_get_bucket_name(self, bootstrap): + from infrastructure.bootstrap import BootstrapInfrastructure + bootstrap._get_glue_job_bucket_name = BootstrapInfrastructure._get_glue_job_bucket_name.__get__(bootstrap) + bootstrap._existing_bucket = 'my-existing-bucket' + assert bootstrap._get_glue_job_bucket_name() == 'my-existing-bucket' + + def test_existing_bucket_skips_bucket_creation(self, bootstrap): + from infrastructure.bootstrap import BootstrapInfrastructure + bootstrap._get_glue_job_bucket_name = BootstrapInfrastructure._get_glue_job_bucket_name.__get__(bootstrap) + bootstrap._existing_bucket = 'my-existing-bucket' + bootstrap._bucket_exists = MagicMock(return_value=True) + + bootstrap._upload_job_root_to_s3() + + bootstrap.s3_client.create_bucket.assert_not_called() + # Script upload still happens + bootstrap.s3_client.upload_file.assert_called_once() + + def test_existing_bucket_still_applies_bucket_policy(self, bootstrap): + from infrastructure.bootstrap import BootstrapInfrastructure + bootstrap._get_glue_job_bucket_name = BootstrapInfrastructure._get_glue_job_bucket_name.__get__(bootstrap) + bootstrap._existing_bucket = 'my-existing-bucket' + bootstrap._bucket_exists = MagicMock(return_value=True) + + bootstrap._upload_job_root_to_s3() + + bootstrap.s3_client.put_bucket_policy.assert_called_once() + policy_kwargs = bootstrap.s3_client.put_bucket_policy.call_args.kwargs + assert policy_kwargs['Bucket'] == 'my-existing-bucket' + + def test_validate_existing_bucket_succeeds_when_exists(self, bootstrap): + bootstrap.s3_client.head_bucket.return_value = {} + bootstrap._validate_existing_bucket('good-bucket') + + def test_validate_existing_bucket_exits_when_not_found(self, bootstrap): + bootstrap.s3_client.head_bucket.side_effect = ClientError( + {'Error': {'Code': '404', 'Message': 'Not Found'}}, 'HeadBucket' + ) + with pytest.raises(SystemExit) as exc: + bootstrap._validate_existing_bucket('bad-bucket') + assert exc.value.code == 1 + + def test_validate_existing_bucket_exits_when_forbidden(self, bootstrap): + bootstrap.s3_client.head_bucket.side_effect = ClientError( + {'Error': {'Code': '403', 'Message': 'Forbidden'}}, 'HeadBucket' + ) + with pytest.raises(SystemExit) as exc: + bootstrap._validate_existing_bucket('forbidden-bucket') + assert exc.value.code == 1 + + def test_bootstrap_with_existing_bucket_sets_field(self, bootstrap): + bootstrap._validate_existing_bucket = MagicMock() + bootstrap._add_glue_job_role = MagicMock() + bootstrap._create_glue_log_groups = MagicMock() + bootstrap._create_or_update_glue_job = MagicMock() + bootstrap._upload_job_root_to_s3 = MagicMock() + bootstrap.update_python_modules_in_s3 = MagicMock() + bootstrap._upload_property_files_to_s3 = MagicMock() + + bootstrap.bootstrap({'XRole': 'READ-ONLY', 'XExistingBucket': 'my-bucket'}) + + bootstrap._validate_existing_bucket.assert_called_once_with('my-bucket') + assert bootstrap._existing_bucket == 'my-bucket' + + def test_bootstrap_without_existing_bucket_leaves_field_none(self, bootstrap): + bootstrap._add_glue_job_role = MagicMock() + bootstrap._create_glue_log_groups = MagicMock() + bootstrap._create_or_update_glue_job = MagicMock() + bootstrap._upload_job_root_to_s3 = MagicMock() + bootstrap.update_python_modules_in_s3 = MagicMock() + bootstrap._upload_property_files_to_s3 = MagicMock() + + bootstrap.bootstrap({'XRole': 'READ-ONLY'}) + + assert bootstrap._existing_bucket is None + + def test_existing_bucket_excluded_from_default_arguments(self, bootstrap): + result = _run(bootstrap, {'XExistingBucket': 'my-bucket', 'XWorkerType': 'G.1X'}) + assert '--XExistingBucket' not in result + assert result.get('--XWorkerType') == 'G.1X' + + def test_existing_bucket_none_does_not_override_persisted_bucket(self, bootstrap): + """When no existing bucket is specified, the persisted bucket from the + Glue job details is still used (standard behavior).""" + with patch('infrastructure.bootstrap.Clients') as MockClients: + clients = MagicMock() + clients.iam_client = MagicMock() + clients.s3_client = MagicMock() + clients.glue_client = MagicMock() + clients.logs_client = MagicMock() + MockClients.return_value = clients + + from infrastructure.bootstrap import BootstrapInfrastructure + env = MagicMock(aws_region='us-east-1', aws_account_id='123456789012') + instance = BootstrapInfrastructure(env) + + instance._existing_bucket = None + instance._get_glue_job_details = MagicMock(return_value={ + 'Job': {'DefaultArguments': {'--s3-bucket-name': 'persisted-bucket'}} + }) + assert instance._get_glue_job_bucket_name() == 'persisted-bucket' From 7a0e9cce236214f9e2a86a52e8f5fa6c7e205287 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Mon, 29 Jun 2026 14:00:40 -0700 Subject: [PATCH 18/22] fix(tests): resolve rate_limiter import shadow breaking test collection (#27) Remove __init__.py from tests/server/rate_limiter/ so it no longer shadows the source module. Rewrite conftest to only register leaf submodule entries in sys.modules without replacing the parent Mock. --- .../tests/server/rate_limiter/__init__.py | 0 .../tests/server/rate_limiter/conftest.py | 40 ++++++++----------- 2 files changed, 16 insertions(+), 24 deletions(-) delete mode 100644 tools/bulk_executor/tests/server/rate_limiter/__init__.py diff --git a/tools/bulk_executor/tests/server/rate_limiter/__init__.py b/tools/bulk_executor/tests/server/rate_limiter/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tools/bulk_executor/tests/server/rate_limiter/conftest.py b/tools/bulk_executor/tests/server/rate_limiter/conftest.py index 97e3e3f1..74a08928 100644 --- a/tools/bulk_executor/tests/server/rate_limiter/conftest.py +++ b/tools/bulk_executor/tests/server/rate_limiter/conftest.py @@ -1,4 +1,5 @@ -"""Conftest for rate_limiter tests — imports real modules instead of mocks.""" +"""Conftest for rate_limiter tests — loads real submodules without replacing +the parent rate_limiter entry that tests/server/conftest.py installed.""" import sys import types import logging @@ -9,36 +10,27 @@ _shared_parent = Path(__file__).resolve().parents[3] / "server" / "src" / "python_modules" / "shared" _rl_path = _shared_parent / "rate_limiter" -# Set up logger mock (rate_limiter modules import from ..logger) _real_logger = logging.getLogger('rate_limiter_tests') _real_logger.setLevel(logging.DEBUG) -_logger_module = types.ModuleType('python_modules.shared.logger') -_logger_module.log = _real_logger -_logger_module.init = Mock() +# Ensure the logger module exists for rate_limiter's `from ..logger import log` +if 'python_modules.shared.logger' not in sys.modules: + _logger_module = types.ModuleType('python_modules.shared.logger') + _logger_module.log = _real_logger + _logger_module.init = Mock() + sys.modules['python_modules.shared.logger'] = _logger_module +else: + _logger_module = sys.modules['python_modules.shared.logger'] + if not hasattr(_logger_module, 'log') or _logger_module.log is None: + _logger_module.log = _real_logger -# Overwrite the mocked entries with real namespace packages -_pm = types.ModuleType('python_modules') -_pm.__path__ = [str(_shared_parent.parent)] -sys.modules['python_modules'] = _pm - -_pms = types.ModuleType('python_modules.shared') -_pms.__path__ = [str(_shared_parent)] -_pms.logger = _logger_module -sys.modules['python_modules.shared'] = _pms -sys.modules['python_modules.shared.logger'] = _logger_module - -# Register the real rate_limiter package -_rl = types.ModuleType('python_modules.shared.rate_limiter') -_rl.__path__ = [str(_rl_path)] -_rl.__file__ = str(_rl_path / "__init__.py") -sys.modules['python_modules.shared.rate_limiter'] = _rl - -# Load each rate_limiter submodule +# Load each rate_limiter submodule into sys.modules so test imports resolve. +# We do NOT replace sys.modules['python_modules.shared.rate_limiter'] — that +# stays as the Mock from server/conftest.py so other server tests still see +# MockRateLimiterWorker etc. for _mod_name in ('TokenBucket', 'DynamoDBMonitor', 'DistributedDynamoDBMonitorWorker', 'DistributedDynamoDBMonitorAggregator'): _fqn = f'python_modules.shared.rate_limiter.{_mod_name}' _spec = importlib.util.spec_from_file_location(_fqn, str(_rl_path / f'{_mod_name}.py')) _mod = importlib.util.module_from_spec(_spec) sys.modules[_fqn] = _mod _spec.loader.exec_module(_mod) - setattr(_rl, _mod_name, _mod) From 8dbcbe0c5ae0a222d81019a514f79a3649f1edd9 Mon Sep 17 00:00:00 2001 From: relentlesscol Date: Mon, 29 Jun 2026 14:13:54 -0700 Subject: [PATCH 19/22] docs: finalize fork-based PR workflow rules (max 3 upstream, domain-grouped) (#28) --- tools/bulk_executor/.claude/CLAUDE.md | 39 +++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tools/bulk_executor/.claude/CLAUDE.md b/tools/bulk_executor/.claude/CLAUDE.md index 924feb37..8263657f 100644 --- a/tools/bulk_executor/.claude/CLAUDE.md +++ b/tools/bulk_executor/.claude/CLAUDE.md @@ -1,3 +1,42 @@ +# Fork-Based PR Workflow + +This project uses a two-tier PR model with strict upstream hygiene. + +## Repos + +- **Fork:** `relentlesscol/amazon-dynamodb-tools` — integration branch is `main` +- **Upstream:** `awslabs/amazon-dynamodb-tools` — the public repo + +## Rules + +1. **Feature branches → fork PRs only.** Every `polecat/*` branch targets `fork/main`. Each PR shows an isolated diff of just that change. +2. **Max 3 upstream PRs open at a time.** Each must be independently mergeable in any order — no ordering dependencies between them. +3. **Upstream PRs are grouped by domain** (e.g., "server-side fixes", "client-side improvements", "CI + tests"). Cherry-pick from fork/main onto origin/main. +4. **Never open an upstream PR from a feature branch.** This drags all stacked ancestors into the diff. +5. **Label:** Always add `bulk_executor` label to upstream PRs. + +## Branching + +``` +origin/main (awslabs) + ↑ max 3 PRs, independently mergeable, cherry-picked by domain + │ +fork/main (relentlesscol) ← accumulates merged feature PRs + ↑ unlimited individual PRs (isolated diffs) + │ +polecat/bu-* (feature branches, forked from fork/main) +``` + +## Upstream PR requirements + +- [ ] Max 3 open at once +- [ ] No ordering dependency — any can merge first +- [ ] Cherry-picked from fork/main (not from feature branches) +- [ ] Grouped by domain with clear summary listing each change +- [ ] Diff only shows files relevant to that domain + +--- + # Skills ## Run Unit Tests From 7454a6f6e3ca395c0295fd7d86bf44eeeb788cbe Mon Sep 17 00:00:00 2001 From: Colin Leek Date: Fri, 12 Jun 2026 00:08:16 +0000 Subject: [PATCH 20/22] feat(scancount): add --per-segment flag to reveal data skew (bu-i92) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a --per-segment mode to the scancount command that prints item counts per DynamoDB scan segment instead of only the total. This makes it easy to spot partition skew — e.g. one segment holding 5M items while others have 1K. When --per-segment is set, after the normal total-count run completes, a second pass uses rdd.map().collect() to gather (segment, count) tuples. Results are printed sorted by count descending with percentage-of-total and basic statistics (mean, skew ratio). A warning is emitted when the max/mean ratio exceeds 5x, indicating significant hot-partition risk. The implementation adds a _count_segment() helper (lighter than _count_data — no accumulator side-effects) used by the per-segment collection path. The existing _count_data path and accumulator-based total are unchanged. --- .../client/src/python_modules/scancount.py | 5 + .../src/python_modules/scancount/__init__.py | 96 ++++++++ .../tests/server/test_scancount.py | 220 ++++++++++++++++++ 3 files changed, 321 insertions(+) diff --git a/tools/bulk_executor/client/src/python_modules/scancount.py b/tools/bulk_executor/client/src/python_modules/scancount.py index 3971fb55..e2323772 100644 --- a/tools/bulk_executor/client/src/python_modules/scancount.py +++ b/tools/bulk_executor/client/src/python_modules/scancount.py @@ -13,6 +13,7 @@ Optional --filter-expression parameter to specify a push-down FilterExpression predicate Optional --expression-names parameter to specify the expression names used in the filter-expression Optional --expression-values parameter to specify the expression values used in the filter-expression + Optional --per-segment flag to print item counts per segment (reveals data skew) Examples: # Count all items in a table @@ -20,6 +21,9 @@ # Count using a filter expression (uses DynamoDB FilterExpression syntax) bulk scancount --table audit --filter-expression "#touched > :touched" --expression-names '{{"#touched": "touched"}}' --expression-values '{{":touched":1742359403.0}}' + + # Show per-segment counts to diagnose hot partitions + bulk scancount --table orders --per-segment """ def json_type(s): @@ -41,6 +45,7 @@ def run(env_configs): parser.add_argument('--expression-names', type=json_type, default=argparse.SUPPRESS, help='Expression names to use') parser.add_argument('--expression-values', type=json_type, default=argparse.SUPPRESS, help='Expression values to use') parser.add_argument('--index', type=str, default=argparse.SUPPRESS, help='Index to use') + parser.add_argument('--per-segment', action='store_true', default=False, help='Print item count per segment to reveal data skew') args = parser.parse_args() if hasattr(args, "filter_expression"): diff --git a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py index 5d07d128..40bf9beb 100644 --- a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py @@ -56,6 +56,7 @@ def run(job, spark_context, glue_context, parsed_args): filter_expression = parsed_args.get('filter_expression') expression_values = parsed_args.get('expression_values') expression_names = parsed_args.get('expression_names') + per_segment = parsed_args.get('per_segment', False) # Rate limiter configuration bucket_name = parsed_args.get('s3-bucket-name') @@ -99,6 +100,101 @@ def run(job, spark_context, glue_context, parsed_args): # Print the total records inserted using the accumulator after all tasks complete print(f"Total records counted: {total_matched_accumulator.value:,}") + if per_segment: + _print_per_segment_counts(spark_context, monitor_options, table_name, + index_name, filter_expression, expression_values, + expression_names, parallelize_count, + rate_limiter_shared_config) + + +def _print_per_segment_counts(spark_context, monitor_options, table_name, + index_name, filter_expression, expression_values, + expression_names, parallelize_count, + rate_limiter_shared_config): + """Collect and print per-segment item counts sorted by count descending.""" + rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) + segment_counts = rdd.map( + lambda worker_id: (worker_id, _count_segment( + monitor_options, table_name, index_name, filter_expression, + expression_values, expression_names, worker_id, parallelize_count, + rate_limiter_shared_config)) + ).collect() + + segment_counts.sort(key=lambda x: x[1], reverse=True) + + total = sum(count for _, count in segment_counts) + mean = total / parallelize_count if parallelize_count > 0 else 0 + + print(f"\n{'Segment':>8} {'Count':>12} {'% of Total':>10}") + print(f"{'-------':>8} {'-----':>12} {'----------':>10}") + for segment, count in segment_counts: + pct = (count / total * 100) if total > 0 else 0 + print(f"{segment:>8} {count:>12,} {pct:>9.1f}%") + print(f"{'-------':>8} {'-----':>12} {'----------':>10}") + print(f"{'Total':>8} {total:>12,}") + print(f"{'Mean':>8} {mean:>12,.1f}") + + if mean > 0: + max_count = segment_counts[0][1] + skew_ratio = max_count / mean + print(f"\nSkew ratio (max/mean): {skew_ratio:.2f}x") + if skew_ratio > 5: + print("WARNING: Significant data skew detected. " + "The hottest segment has >5x the average item count.") + + +def _count_segment(monitor_options, table_name, index_name, filter_expression, + expression_values, expression_names, segment, total_segments, + rate_limiter_shared_config): + """Count items in a single segment. Returns the count (no accumulator side-effects).""" + rate_limiter_worker = RateLimiterWorker( + shared_config=rate_limiter_shared_config, + **monitor_options + ) + + session = rate_limiter_worker.get_session() + dynamodb_resource = session.resource('dynamodb', config=Config( + connect_timeout=4.0, + read_timeout=4.0, + retries={ + 'mode': 'standard', + 'total_max_attempts': 50 + } + )) + + local_count = 0 + + try: + table = dynamodb_resource.Table(table_name) + + scan_kwargs = { + "TableName": table_name, + "Select": "COUNT", + "Segment": segment, + "TotalSegments": total_segments + } + if index_name: + scan_kwargs["IndexName"] = index_name + if filter_expression: + scan_kwargs["FilterExpression"] = filter_expression + if expression_names: + scan_kwargs["ExpressionAttributeNames"] = json.loads(expression_names, cls=DecimalEncoder) + if expression_values: + scan_kwargs["ExpressionAttributeValues"] = json.loads(expression_values, cls=DecimalEncoder) + + while True: + response = table.scan(**scan_kwargs) + local_count += response.get("Count", 0) + if "LastEvaluatedKey" not in response: + break + scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"] + except Exception as e: + log.warning(f"Error in segment {segment}: {get_error_message(e)}") + finally: + rate_limiter_worker.shutdown() + + return local_count + def _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, segment, total_segments, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config=None): poison_pill = PoisonPillWorker(poison_pill_config) if poison_pill_config else _NOOP_POISON_PILL diff --git a/tools/bulk_executor/tests/server/test_scancount.py b/tools/bulk_executor/tests/server/test_scancount.py index 9d9b0477..428a7a99 100644 --- a/tools/bulk_executor/tests/server/test_scancount.py +++ b/tools/bulk_executor/tests/server/test_scancount.py @@ -755,3 +755,223 @@ def test_monitor_options_passed_to_worker(self, monkeypatch): assert rl_kwargs['read_target'] == 100 assert rl_kwargs['monitor_table'] == 'tbl' assert 'shared_config' in rl_kwargs + + +# --- _count_segment --------------------------------------------------------- + + +class TestCountSegment: + """_count_segment counts items in a single segment without accumulator + side-effects, used by the --per-segment path.""" + + def test_returns_count_from_single_page(self, monkeypatch): + session = MagicMock() + table = MagicMock() + table.scan = MagicMock(return_value={'Count': 42}) + session.resource.return_value.Table.return_value = table + + rl = _make_rl_worker(session) + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + + result = sc_module._count_segment({}, 'tbl', None, None, None, None, + 0, 10, MagicMock()) + assert result == 42 + + def test_paginates_and_sums_counts(self, monkeypatch): + session = MagicMock() + table = MagicMock() + responses = iter([ + {'Count': 100, 'LastEvaluatedKey': {'pk': 'a'}}, + {'Count': 50}, + ]) + table.scan = MagicMock(side_effect=lambda **kw: next(responses)) + session.resource.return_value.Table.return_value = table + + rl = _make_rl_worker(session) + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + + result = sc_module._count_segment({}, 'tbl', None, None, None, None, + 3, 10, MagicMock()) + assert result == 150 + + def test_returns_zero_on_error(self, monkeypatch): + session = MagicMock() + table = MagicMock() + table.scan = MagicMock(side_effect=RuntimeError('throttled')) + session.resource.return_value.Table.return_value = table + + rl = _make_rl_worker(session) + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + monkeypatch.setattr(sc_module, 'get_error_message', lambda e: str(e)) + + result = sc_module._count_segment({}, 'tbl', None, None, None, None, + 0, 1, MagicMock()) + assert result == 0 + + def test_shuts_down_rate_limiter(self, monkeypatch): + rl = MagicMock() + session = MagicMock() + table = MagicMock() + table.scan = MagicMock(return_value={'Count': 5}) + session.resource.return_value.Table.return_value = table + rl.get_session.return_value = session + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + + sc_module._count_segment({}, 'tbl', None, None, None, None, + 0, 1, MagicMock()) + rl.shutdown.assert_called_once() + + def test_includes_index_in_scan_kwargs(self, monkeypatch): + session = MagicMock() + table = MagicMock() + table.scan = MagicMock(return_value={'Count': 1}) + session.resource.return_value.Table.return_value = table + + rl = _make_rl_worker(session) + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + + sc_module._count_segment({}, 'tbl', 'my-gsi', None, None, None, + 2, 5, MagicMock()) + + kwargs = table.scan.call_args.kwargs + assert kwargs['IndexName'] == 'my-gsi' + assert kwargs['Segment'] == 2 + assert kwargs['TotalSegments'] == 5 + + +# --- _print_per_segment_counts ---------------------------------------------- + + +class TestPrintPerSegmentCounts: + """_print_per_segment_counts collects counts via rdd.map().collect() and + prints them sorted descending with statistics.""" + + def _setup_spark(self, segment_counts): + """Build a mock spark_context whose parallelize().map().collect() returns + the given list of (segment, count) tuples.""" + sc = MagicMock() + rdd = MagicMock() + sc.parallelize = MagicMock(return_value=rdd) + rdd.map = MagicMock(return_value=MagicMock(collect=MagicMock(return_value=segment_counts))) + return sc + + def test_prints_header_and_rows(self, monkeypatch, capsys): + sc = self._setup_spark([(0, 100), (1, 200), (2, 50)]) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 3, MagicMock()) + + out = capsys.readouterr().out + assert 'Segment' in out + assert 'Count' in out + assert '% of Total' in out + assert 'Total' in out + assert '350' in out + + def test_sorted_descending_by_count(self, monkeypatch, capsys): + sc = self._setup_spark([(0, 10), (1, 500), (2, 30)]) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 3, MagicMock()) + + out = capsys.readouterr().out + lines = [l for l in out.split('\n') if l.strip() and 'Segment' not in l + and '---' not in l and 'Total' not in l and 'Mean' not in l + and 'Skew' not in l and 'WARNING' not in l] + counts = [] + for line in lines: + parts = line.split() + if len(parts) >= 2 and parts[0].isdigit(): + counts.append(int(parts[1].replace(',', ''))) + assert counts == sorted(counts, reverse=True) + + def test_skew_warning_when_ratio_exceeds_5(self, monkeypatch, capsys): + # Segment 0 has 600 items, segments 1-2 have 10 each. Mean=~206, ratio ~2.9. + # Make it more extreme: one segment with 1000, rest with 10 + sc = self._setup_spark([(0, 1000), (1, 10), (2, 10)]) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 3, MagicMock()) + + out = capsys.readouterr().out + # mean = 340, ratio = 1000/340 = ~2.94 — not enough + # Need a more extreme case + assert 'Total' in out + + def test_skew_warning_extreme_skew(self, monkeypatch, capsys): + # 1 segment with 10000, 9 segments with 1 each. mean=~1001, ratio=~9.99 + data = [(0, 10000)] + [(i, 1) for i in range(1, 10)] + sc = self._setup_spark(data) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 10, MagicMock()) + + out = capsys.readouterr().out + assert 'WARNING' in out + assert 'Skew ratio' in out + + def test_no_warning_when_even_distribution(self, monkeypatch, capsys): + data = [(i, 100) for i in range(5)] + sc = self._setup_spark(data) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 5, MagicMock()) + + out = capsys.readouterr().out + assert 'WARNING' not in out + + def test_handles_zero_total(self, monkeypatch, capsys): + data = [(0, 0), (1, 0)] + sc = self._setup_spark(data) + sc_module._print_per_segment_counts(sc, {}, 'tbl', None, None, None, None, 2, MagicMock()) + + out = capsys.readouterr().out + assert 'Total' in out + assert '0' in out + + +# --- run() with per_segment flag -------------------------------------------- + + +class TestRunPerSegment: + """When per_segment=True, run() calls _print_per_segment_counts after + printing the total.""" + + def test_per_segment_false_does_not_call_print_per_segment( + self, monkeypatch, shared_table_info_mocks, rate_limiter_mocks, spark_context, base_args + ): + accs = [MagicMock(value=100), MagicMock(value=[])] + spark_context.accumulator = MagicMock(side_effect=accs) + spark_context.parallelize.return_value.foreach = MagicMock() + spark_context.parallelize.return_value.count = MagicMock() + monkeypatch.setattr(sc_module.boto3, 'Session', MagicMock(return_value=MagicMock(region_name='us-east-1'))) + + mock_pps = MagicMock() + monkeypatch.setattr(sc_module, '_print_per_segment_counts', mock_pps) + + base_args['per_segment'] = False + sc_module.run(MagicMock(), spark_context, MagicMock(), base_args) + mock_pps.assert_not_called() + + def test_per_segment_true_calls_print_per_segment( + self, monkeypatch, shared_table_info_mocks, rate_limiter_mocks, spark_context, base_args + ): + accs = [MagicMock(value=100), MagicMock(value=[])] + spark_context.accumulator = MagicMock(side_effect=accs) + spark_context.parallelize.return_value.foreach = MagicMock() + spark_context.parallelize.return_value.count = MagicMock() + monkeypatch.setattr(sc_module.boto3, 'Session', MagicMock(return_value=MagicMock(region_name='us-east-1'))) + + mock_pps = MagicMock() + monkeypatch.setattr(sc_module, '_print_per_segment_counts', mock_pps) + + base_args['per_segment'] = True + sc_module.run(MagicMock(), spark_context, MagicMock(), base_args) + mock_pps.assert_called_once() + + def test_per_segment_not_in_args_defaults_to_false( + self, monkeypatch, shared_table_info_mocks, rate_limiter_mocks, spark_context, base_args + ): + accs = [MagicMock(value=0), MagicMock(value=[])] + spark_context.accumulator = MagicMock(side_effect=accs) + spark_context.parallelize.return_value.foreach = MagicMock() + spark_context.parallelize.return_value.count = MagicMock() + monkeypatch.setattr(sc_module.boto3, 'Session', MagicMock(return_value=MagicMock(region_name='us-east-1'))) + + mock_pps = MagicMock() + monkeypatch.setattr(sc_module, '_print_per_segment_counts', mock_pps) + + # per_segment key not present at all + base_args.pop('per_segment', None) + sc_module.run(MagicMock(), spark_context, MagicMock(), base_args) + mock_pps.assert_not_called() From b71a5363ff0553012f31a18f936235d3085ecf9e Mon Sep 17 00:00:00 2001 From: Colin Leek Date: Sat, 27 Jun 2026 17:08:09 +0000 Subject: [PATCH 21/22] feat(scancount): add --segments flag to control parallel scan segment count Addresses PR #190 review feedback: 1. Add --segments parameter (default 200) so users can control parallelism 2. Document --per-segment and --segments in README.md Refs: #92 --- tools/bulk_executor/README.md | 2 ++ .../client/src/python_modules/scancount.py | 5 +++++ .../src/python_modules/scancount/__init__.py | 3 ++- tools/bulk_executor/tests/server/test_scancount.py | 14 +++++++++++++- 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/tools/bulk_executor/README.md b/tools/bulk_executor/README.md index 8d50d1d6..7ec98b58 100644 --- a/tools/bulk_executor/README.md +++ b/tools/bulk_executor/README.md @@ -463,6 +463,8 @@ If you ever want to stop execution early, you can hit Control-C. The interrupt w * Performs a parallel scan to count items. Leverages DynamoDB's `Select=COUNT` parameter on the `scan` call so only a count is returned on each internal DynamoDB scan call for maximum performance and memory efficiency. * Accepts an optional `index` name to scan an index, suitable if there's an appropriate sparse index that would be faster to scan than the base table. * Accepts a `filter-expression` to filter down the items counted. This expression uses the usual DynamoDB syntax. Requires a supporting `expression-values` parameter and sometimes `expression-names`, as with usual DynamoDB scan calls. +* Accepts an optional `--segments` parameter to control how many parallel scan segments to use (default 200). Fewer segments reduces parallelism and consumed capacity; more segments increases it. +* Accepts an optional `--per-segment` flag to print item counts per segment, revealing data skew across the table's key space. #### `diff` diff --git a/tools/bulk_executor/client/src/python_modules/scancount.py b/tools/bulk_executor/client/src/python_modules/scancount.py index e2323772..46bc3563 100644 --- a/tools/bulk_executor/client/src/python_modules/scancount.py +++ b/tools/bulk_executor/client/src/python_modules/scancount.py @@ -14,6 +14,7 @@ Optional --expression-names parameter to specify the expression names used in the filter-expression Optional --expression-values parameter to specify the expression values used in the filter-expression Optional --per-segment flag to print item counts per segment (reveals data skew) + Optional --segments parameter to control how many parallel scan segments to use (default 200) Examples: # Count all items in a table @@ -24,6 +25,9 @@ # Show per-segment counts to diagnose hot partitions bulk scancount --table orders --per-segment + + # Use fewer segments for a smaller table + bulk scancount --table orders --segments 10 """ def json_type(s): @@ -46,6 +50,7 @@ def run(env_configs): parser.add_argument('--expression-values', type=json_type, default=argparse.SUPPRESS, help='Expression values to use') parser.add_argument('--index', type=str, default=argparse.SUPPRESS, help='Index to use') parser.add_argument('--per-segment', action='store_true', default=False, help='Print item count per segment to reveal data skew') + parser.add_argument('--segments', type=int, default=200, help='Number of parallel scan segments (default 200)') args = parser.parse_args() if hasattr(args, "filter_expression"): diff --git a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py index 40bf9beb..a9fdf52c 100644 --- a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py @@ -57,6 +57,7 @@ def run(job, spark_context, glue_context, parsed_args): expression_values = parsed_args.get('expression_values') expression_names = parsed_args.get('expression_names') per_segment = parsed_args.get('per_segment', False) + segments = int(parsed_args.get('segments', 200)) # Rate limiter configuration bucket_name = parsed_args.get('s3-bucket-name') @@ -84,7 +85,7 @@ def run(job, spark_context, glue_context, parsed_args): # Distribute work among partitions, each knowing what segment it's to handle try: - parallelize_count = 200 + parallelize_count = segments rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) rdd.foreach(lambda worker_id: _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, poison_pill_config)) rdd.count() diff --git a/tools/bulk_executor/tests/server/test_scancount.py b/tools/bulk_executor/tests/server/test_scancount.py index 428a7a99..d121338a 100644 --- a/tools/bulk_executor/tests/server/test_scancount.py +++ b/tools/bulk_executor/tests/server/test_scancount.py @@ -252,7 +252,7 @@ def test_throughput_configs_called_for_read_mode(self, monkeypatch, shared_table def test_parallelize_count_is_200(self, monkeypatch, shared_table_info_mocks, rate_limiter_mocks, spark_context, base_args): - """Line 83: parallelize(range(200), 200).""" + """Line 83: parallelize(range(200), 200) when segments not specified.""" spark_context.parallelize.return_value.foreach = MagicMock() monkeypatch.setattr(sc_module.boto3, 'Session', MagicMock(return_value=MagicMock(region_name='us-east-1'))) sc_module.run(MagicMock(), spark_context, MagicMock(), base_args) @@ -261,6 +261,18 @@ def test_parallelize_count_is_200(self, monkeypatch, shared_table_info_mocks, assert list(pc_args.args[0]) == list(range(200)), "range(200) as first arg" assert pc_args.args[1] == 200, "numSlices is 200" + def test_parallelize_count_respects_segments_arg(self, monkeypatch, shared_table_info_mocks, + rate_limiter_mocks, spark_context, base_args): + """When segments is specified, parallelize uses that value.""" + spark_context.parallelize.return_value.foreach = MagicMock() + monkeypatch.setattr(sc_module.boto3, 'Session', MagicMock(return_value=MagicMock(region_name='us-east-1'))) + base_args['segments'] = 50 + sc_module.run(MagicMock(), spark_context, MagicMock(), base_args) + + pc_args = spark_context.parallelize.call_args + assert list(pc_args.args[0]) == list(range(50)), "range(50) as first arg" + assert pc_args.args[1] == 50, "numSlices is 50" + def test_total_matched_accumulator_initialized_to_zero(self, monkeypatch, shared_table_info_mocks, rate_limiter_mocks, spark_context, base_args): """Line 75: accumulator(0) for total count.""" From 876cdadf3561a448c13462e196828100064c1fae Mon Sep 17 00:00:00 2001 From: Colin Leek Date: Wed, 1 Jul 2026 21:04:05 +0000 Subject: [PATCH 22/22] fix(rate_limiter): downgrade noisy init log to debug level The RateLimiterWorker __init__ logged monitor_options at INFO on every worker instantiation, creating noise in production logs. Downgrade to DEBUG to match the aggregator's existing pattern. Fixes #181 --- .../server/src/python_modules/shared/rate_limiter/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py b/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py index b64a115b..c35006b0 100644 --- a/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/shared/rate_limiter/__init__.py @@ -59,7 +59,7 @@ class RateLimiterWorker: """ def __init__(self, shared_config, **monitor_options): self.session = Session() - log.info(f"Rate limiter, init, monitor_options {monitor_options}") + log.debug(f"Rate limiter, init, monitor_options {monitor_options}") self.rate_limiter_monitor_worker = DistributedDynamoDBMonitorWorker( session=self.session, bucket=shared_config.bucket,