sync-dynamodb.sh

# Wrap logic in main function to prevent 'exit' from killing shell if sourced
main() {
  # Load .env file if it exists
  if [[ -f .env ]]; then
    set -a
    source .env
    set +a
  fi

  local LOCAL_ENDPOINT="${LOCAL_ENDPOINT:-http://localhost:4567}"
  local REGION="${REGION:-}"
  local TABLES_ARRAY=()
  
  if [[ -n "${TABLES:-}" ]]; then
    if [[ -n "${ZSH_VERSION:-}" ]]; then
      # Zsh array reading
      read -r -A TABLES_ARRAY <<< "$TABLES"
    else
      # Bash array reading
      IFS=' ' read -r -a TABLES_ARRAY <<< "$TABLES"
    fi
  fi

  # Helper for usage
  show_usage() {
    echo "Usage: $0 [-t \"Table1 Table2\"] [-r us-west-2] [-e http://localhost:4567]"
    echo ""
    echo "Options:"
    echo "  -t, --tables     Required (if not in .env). Space-separated list of tables to sync."
    echo "  -r, --region     Required (if not in .env). AWS Region (e.g., us-west-2)."
    echo "  -e, --endpoint   Optional. Local DynamoDB endpoint (default: ${LOCAL_ENDPOINT})."
    echo "  -h, --help       Show this help message."
    return 1
  }

  # Parse arguments
  local ARGS=("$@")
  local idx=0
  while [[ $idx -lt ${#ARGS[@]} ]]; do
    local key="${ARGS[$idx]}"
    case $key in
      -t|--tables)
        idx=$((idx+1))
        if [[ -n "${ZSH_VERSION:-}" ]]; then
          read -r -A TABLES_ARRAY <<< "${ARGS[$idx]}"
        else
          IFS=' ' read -r -a TABLES_ARRAY <<< "${ARGS[$idx]}"
        fi
        ;;
      -r|--region)
        idx=$((idx+1))
        REGION="${ARGS[$idx]}"
        ;;
      -e|--endpoint)
        idx=$((idx+1))
        LOCAL_ENDPOINT="${ARGS[$idx]}"
        ;;
      -h|--help)
        show_usage
        return 1
        ;;
      *)
        # Only error if it looks like a flag, otherwise execution might continue? 
        # But simplistic parsing is safer.
        # echo "Unknown option: $key"
        # show_usage
        # return 1
        ;;
    esac
    idx=$((idx+1))
  done

  local TABLES=("${TABLES_ARRAY[@]}")

  if [[ ${#TABLES[@]} -eq 0 ]]; then
    echo "Error: Tables are required via -t or .env"
    show_usage
    return 1
  fi

  if [[ -z "${REGION}" ]]; then
    echo "Error: Region is required via -r or .env"
    show_usage
    return 1
  fi

  local TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
  local BACKUP_DIR="./backup/${TIMESTAMP}"

  mkdir -p "${BACKUP_DIR}"
  echo "==> Backup directory: ${BACKUP_DIR}"

  # 1. Export production data
  for TABLE in "${TABLES[@]}"; do
    echo "==> Exporting table data: ${TABLE}"
    aws dynamodb scan \
      --table-name "${TABLE}" \
      --region "${REGION}" \
      --output json \
      > "${BACKUP_DIR}/${TABLE}.json" || return 1
  done

  # 2. Export schema (READ ONLY - PROD)
  for TABLE in "${TABLES[@]}"; do
    echo "==> Exporting schema: ${TABLE}"
    local SCHEMA_FILE="${BACKUP_DIR}/${TABLE}-schema.json"

    aws dynamodb describe-table \
      --table-name "${TABLE}" \
      --region "${REGION}" \
      | jq '{
          TableName: .Table.TableName,
          AttributeDefinitions: .Table.AttributeDefinitions,
          KeySchema: .Table.KeySchema,
          BillingMode: "PAY_PER_REQUEST"
        }
        + (if .Table.GlobalSecondaryIndexes then
            {
              GlobalSecondaryIndexes: (
                .Table.GlobalSecondaryIndexes
                | map({
                    IndexName,
                    KeySchema,
                    Projection
                  })
              )
            }
          else {} end)
      ' > "${SCHEMA_FILE}" || return 1
  done

  # SAFETY CHECK: Remove credentials
  cleanup

  # Set dummy credentials for local operations
  export AWS_ACCESS_KEY_ID="dummy"
  export AWS_SECRET_ACCESS_KEY="dummy"
  export AWS_REGION="us-west-2"

  # 3. Recreate local tables
  for TABLE in "${TABLES[@]}"; do
    local SCHEMA_FILE="${BACKUP_DIR}/${TABLE}-schema.json"
    
    if aws dynamodb describe-table \
        --table-name "${TABLE}" \
        --endpoint-url "${LOCAL_ENDPOINT}" >/dev/null 2>&1; then

      echo "==> Deleting local table: ${TABLE}"
      aws dynamodb delete-table \
        --table-name "${TABLE}" \
        --endpoint-url "${LOCAL_ENDPOINT}" > /dev/null || return 1

      aws dynamodb wait table-not-exists \
        --table-name "${TABLE}" \
        --endpoint-url "${LOCAL_ENDPOINT}"
    fi

    echo "==> Creating local table: ${TABLE}"
    aws dynamodb create-table \
      --cli-input-json file://"${SCHEMA_FILE}" \
      --endpoint-url "${LOCAL_ENDPOINT}" > /dev/null || return 1

    aws dynamodb wait table-exists \
      --table-name "${TABLE}" \
      --endpoint-url "${LOCAL_ENDPOINT}"
      
    echo
  done

  # 4. Import data
  for TABLE in "${TABLES[@]}"; do
    echo "==> Importing data into local table: ${TABLE}"
    if [[ -f "${BACKUP_DIR}/${TABLE}.json" ]]; then
       jq -c '.Items[]' "${BACKUP_DIR}/${TABLE}.json" | while read -r ITEM; do
         aws dynamodb put-item \
           --table-name "${TABLE}" \
           --item "${ITEM}" \
           --endpoint-url "${LOCAL_ENDPOINT}"
       done
    fi
  done

  echo
  echo "==> Done. Local DynamoDB is now fully synced with prod."

  # Final cleanup
  cleanup
}

# Cleanup function to unset credentials
cleanup() {
  echo
  # Call aws-env.sh to unset credentials.
  source ./aws-env.sh unauth
  echo
}

# Check if script is sourced
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  echo "Error: This script must be sourced to properly handle credentials."
  echo "Please run: source $0"
  exit 1
fi

# Run main with arguments
main "$@"