Automate On-the-Fly Backup & Restore of Specific YCQL Tables via YBA API

				
					#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOS'
Usage:
  yb-ycql-table-refresh.sh \
    --yba-url https://portal.example.com \
    --api-token <TOKEN> \
    --customer-uuid <CUUID> \
    --source-universe <SRC_UUID> \
    --target-universe <TGT_UUID> \
    --keyspace <ks> \
    --tables t2,t3 \
    --storage-config <SCUUID> \
    [--poll-interval 20] \
    [--sse true|false] [--use-tablespaces true|false] \
    [--pre-drop] [--target-ycql-host <host[:9042]>] \
    [--delete-after] \
    [--target-ycql-user <user>] [--target-ycql-pass <pass>] [--target-ycql-pass-file <file>]

What it does (YBA ≥ 2.29, YCQL):
- Resolves table UUIDs for a subset of YCQL tables (skips index tables).
- Performs a single v2 backup (keyspaceTableList with tableUUIDList).
- Finds the backup via /backups/page.
- (If --pre-drop) Drops target tables AFTER backup succeeds, right before restore.
- Restores only the specified tables to the target universe.
- Displays exact YBA task errors via /tasks/<uuid>/failed (falls back to /subtasks).
- (If --delete-after) Deletes the backup via POST /backups/delete and waits on task.

Secure YCQL auth:
- Prefer env: YB_YCQL_USER / YB_YCQL_PASSWORD
- Or flags: --target-ycql-user / --target-ycql-pass (or --target-ycql-pass-file)
- Script writes a 0600 temp cqlshrc and uses --cqlshrc.

Requires: curl, jq. (ycqlsh required for --pre-drop or target diagnostics)
EOS
}

# ---------- Defaults
POLL_INTERVAL=20
DELETE_AFTER=false
PRE_DROP=false
TARGET_YCQL_HOST=""
SSE=true
USE_TABLESPACES=false

# ---------- YCQL auth (Option A)
TARGET_YCQL_USER="${YB_YCQL_USER:-}"
TARGET_YCQL_PASS=""
TARGET_YCQL_PASS_FILE=""

# ---------- Args
while [[ $# -gt 0 ]]; do
  case "$1" in
    --yba-url) YBA_URL="$2"; shift 2;;
    --api-token) API_TOKEN="$2"; shift 2;;
    --customer-uuid) CUUID="$2"; shift 2;;
    --source-universe) SRC_UUUID="$2"; shift 2;;
    --target-universe) TGT_UUUID="$2"; shift 2;;
    --keyspace) KEYSPACE="$2"; shift 2;;
    --tables) TABLES_CSV="$2"; shift 2;;
    --storage-config) SCUUID="$2"; shift 2;;
    --poll-interval) POLL_INTERVAL="$2"; shift 2;;
    --sse) SSE="$2"; shift 2;;
    --use-tablespaces) USE_TABLESPACES="$2"; shift 2;;
    --delete-after) DELETE_AFTER=true; shift 1;;
    --pre-drop) PRE_DROP=true; shift 1;;
    --target-ycql-host) TARGET_YCQL_HOST="$2"; shift 2;;
    --target-ycql-user) TARGET_YCQL_USER="$2"; shift 2;;
    --target-ycql-pass) TARGET_YCQL_PASS="$2"; shift 2;;
    --target-ycql-pass-file) TARGET_YCQL_PASS_FILE="$2"; shift 2;;
    -h|--help) usage; exit 0;;
    *) echo "Unknown arg: $1" >&2; usage; exit 1;;
  esac
done

# Resolve YCQL password (prefer file > flag > env)
if [[ -n "$TARGET_YCQL_PASS_FILE" ]]; then
  TARGET_YCQL_PASS="$(<"$TARGET_YCQL_PASS_FILE")"
elif [[ -z "$TARGET_YCQL_PASS" && -n "${YB_YCQL_PASSWORD:-}" ]]; then
  TARGET_YCQL_PASS="$YB_YCQL_PASSWORD"
fi

# ---------- Validate
for v in YBA_URL API_TOKEN CUUID SRC_UUUID TGT_UUUID KEYSPACE TABLES_CSV SCUUID; do
  [[ -n "${!v:-}" ]] || { echo "Missing --${v//_/-}" >&2; usage; exit 1; }
done
command -v curl >/dev/null || { echo "curl required" >&2; exit 1; }
command -v jq   >/dev/null || { echo "jq required" >&2; exit 1; }
if $PRE_DROP; then command -v ycqlsh >/dev/null || { echo "ycqlsh required for --pre-drop" >&2; exit 1; }; fi

AUTH="X-AUTH-YW-API-TOKEN: $API_TOKEN"
CT="Content-Type: application/json"

# ---------- Secure ycqlsh wrapper (Option A)
YCQL_RCFILE=""
make_ycqlsh_cmd() {
  local host="$1" port="$2"
  if [[ -z "$TARGET_YCQL_USER" || -z "$TARGET_YCQL_PASS" ]]; then
    echo "ycqlsh $host $port"; return
  fi
  if [[ -z "$YCQL_RCFILE" ]]; then
    YCQL_RCFILE="$(mktemp)"
    chmod 600 "$YCQL_RCFILE"
    cat >"$YCQL_RCFILE" <<EOF
[authentication]
username = $TARGET_YCQL_USER
password = $TARGET_YCQL_PASS
EOF
  fi
  echo "ycqlsh --cqlshrc \"$YCQL_RCFILE\" $host $port"
}
cleanup_ycqlrc() { [[ -n "$YCQL_RCFILE" && -f "$YCQL_RCFILE" ]] && rm -f "$YCQL_RCFILE"; }
trap cleanup_ycqlrc EXIT

# ---------- Error detail helpers
print_failure_details() {
  local task_uuid="$1"
  echo "---- YBA failure details for task $task_uuid ----"
  local FAILED
  FAILED=$(curl -sS -H "$AUTH" "$YBA_URL/api/v1/customers/$CUUID/tasks/$task_uuid/failed" || true)
  if echo "$FAILED" | jq -e . >/dev/null 2>&1; then
    echo "$FAILED" | jq -r '
      .failedSubTasks // []
      | if length==0 then "No failedSubTasks reported."
        else .[] |
          "subTaskType: " + (.subTaskType // "unknown") + "\n" +
          "state      : " + (.subTaskState // "") + "\n" +
          "error      : " + ((.errorString // "") | gsub("\\n"; "\n              ")) + "\n"
        end
    '
  else
    echo "failed-subtasks endpoint not available (or non-JSON)."
  fi
  local SUBS
  SUBS=$(curl -sS -H "$AUTH" "$YBA_URL/api/v1/customers/$CUUID/tasks/$task_uuid/subtasks" || true)
  if echo "$SUBS" | jq -e . >/dev/null 2>&1; then
    echo "$SUBS" | jq -r '
      (if type=="array" then . else (.entities // .data // []) end)
      | map({
          title: (.title // .name // .subTaskGroupType // "subtask"),
          status: (.status // .state // ""),
          err: (.errorString // .details?.errorString // .details?.error?.message // "")
        })
      | map(select(.err != "" or (.status|ascii_downcase)=="failure"))
      | if length==0 then "No errorString fields found in subtasks."
        else .[] | "\(.title): \(.err)\n"
        end
    '
  fi
  local TOP
  TOP=$(curl -sS -H "$AUTH" "$YBA_URL/api/v1/customers/$CUUID/tasks/$task_uuid" || true)
  if echo "$TOP" | jq -e . >/dev/null 2>&1; then
    echo "---- Task summary ----"
    echo "$TOP" | jq '{title,status,percent,details}'
  fi
  echo "-----------------------------------------------"
}

diagnose_target_existing_tables() {
  [[ -n "$TARGET_YCQL_HOST" ]] || return 0
  command -v ycqlsh >/dev/null || return 0
  local ks="$1" tables_csv="$2"
  local host="${TARGET_YCQL_HOST%%:*}"
  local port="${TARGET_YCQL_HOST##*:}"; [[ "$host" == "$port" ]] && port="9042"
  local YCQLSH_CMD; YCQLSH_CMD=$(make_ycqlsh_cmd "$host" "$port")
  echo "---- Target YCQL tables already existing in ${ks} ----"
  eval "$YCQLSH_CMD -e \"SELECT table_name FROM system_schema.tables WHERE keyspace_name='${ks}' AND table_name IN ('${tables_csv//,/','}');\"" 2>/dev/null || true
  echo "------------------------------------------------------"
}

wait_task() {
  local label="$1" task_uuid="$2"
  while true; do
    local st
    st=$(curl -sS -H "$AUTH" "$YBA_URL/api/v1/customers/$CUUID/tasks/$task_uuid" | jq -r '.status // .state // empty')
    echo "   $label: ${st:-unknown}"
    if [[ "$st" == "Success" || "$st" == "Completed" ]]; then
      return 0
    elif [[ "$st" == "Failure" || "$st" == "Failed" ]]; then
      print_failure_details "$task_uuid"
      if [[ "$label" == "restore" ]]; then
        diagnose_target_existing_tables "$KEYSPACE" "$TABLES_CSV"
        echo "Hint: YCQL restore will not overwrite existing tables. Use --pre-drop --target-ycql-host <host[:9042]> to drop before restore."
      fi
      return 1
    fi
    sleep "$POLL_INTERVAL"
  done
}

# ---------- Universe table listing
fetch_tables_json() {
  local url r
  url="$YBA_URL/api/v1/customers/$CUUID/universes/$SRC_UUUID/tables?tableType=YQL_TABLE_TYPE&includeSystemTables=false&excludeIndexTables=true"
  r=$(curl -sS -H "$AUTH" "$url" 2>/dev/null || true)
  if echo "$r" | jq -e . >/dev/null 2>&1; then echo "$r"; return; fi
  url="$YBA_URL/api/v1/customers/$CUUID/universes/$SRC_UUUID/tables?tableType=YQL_TABLE_TYPE&includeSystemTables=false&excludeIndexTables=false"
  r=$(curl -sS -H "$AUTH" "$url" 2>/dev/null || true)
  if echo "$r" | jq -e . >/dev/null 2>&1; then echo "$r"; return; fi
  url="$YBA_URL/api/v1/customers/$CUUID/universes/$SRC_UUUID/tables"
  r=$(curl -sS -H "$AUTH" "$url" 2>/dev/null || true)
  if echo "$r" | jq -e . >/dev/null 2>&1; then echo "$r"; return; fi
  echo "[]"
}

# ---------- Inputs
IFS=',' read -r -a TABLES_ARR <<< "$TABLES_CSV"
TABLES_JSON=$(printf '%s\n' "${TABLES_ARR[@]}" | jq -R . | jq -s .)

# ---------- Resolve table UUIDs (skip indexes)
echo "==> Resolving YCQL table UUIDs for keyspace '$KEYSPACE' and tables: $TABLES_CSV"
ALL_TABLES_JSON=$(fetch_tables_json)
NORM=$(echo "$ALL_TABLES_JSON" | jq '
  (if type=="array" then . else (.data // .entities // []) end)
  | map({
      tableName: (.tableName // .name // ""),
      tableUUID: (.tableUUID // .uuid // .id // .tableID // ""),
      tableType: (.tableType // .type // ""),
      isIndex: (.isIndex // .isIndexTable // false)
    })
  | map(select(.tableType == "YQL_TABLE_TYPE"))
  | map(select((.isIndex | not) and (.tableName | test("_idx$") | not)))
')
TBL_UUIDS_JSON=$(jq -n --argjson want "$TABLES_JSON" --argjson all "$NORM" '
  $all | map(select(.tableName as $n | ($want | index($n)))) | map(.tableUUID)
')
echo "==> Resolved table UUIDs:"
jq -n --argjson want "$TABLES_JSON" --argjson all "$NORM" '
  $all | map(select(.tableName as $n | ($want | index($n)))) | map({tableName, tableUUID})
' | jq .
if [[ -z "$TBL_UUIDS_JSON" || "$TBL_UUIDS_JSON" == "[]" ]]; then
  echo "Could not resolve YCQL table UUIDs for: $TABLES_CSV" >&2
  echo "$NORM" | jq . >&2
  exit 50
fi
COUNT_WANT=$(echo "$TABLES_JSON" | jq 'length')
COUNT_UUID=$(echo "$TBL_UUIDS_JSON" | jq 'length')
if [[ "$COUNT_WANT" -ne "$COUNT_UUID" ]]; then
  echo "Resolved $COUNT_UUID UUID(s) for $COUNT_WANT requested table(s) — mismatch. Aborting." >&2
  echo "Names:";  echo "$TABLES_JSON"    | jq .
  echo "UUIDs:";  echo "$TBL_UUIDS_JSON" | jq .
  exit 51
fi

# ---------- 1) Create ONE subset backup (names + UUIDs)
echo "==> Creating FULL YCQL backup for keyspace '$KEYSPACE' tables: $TABLES_CSV"
BACKUP_PAYLOAD=$(jq -n \
  --arg u "$SRC_UUUID" --arg sc "$SCUUID" --arg ks "$KEYSPACE" \
  --argjson tl "$TABLES_JSON" --argjson tul "$TBL_UUIDS_JSON" \
  --arg cu "$CUUID" --argjson sse "$SSE" --argjson ut "$USE_TABLESPACES" '
  {
    backupType:"YQL_TABLE_TYPE", customerUUID:$cu, sse:$sse,
    storageConfigUUID:$sc, universeUUID:$u,
    tableByTableBackup:false, useTablespaces:$ut,
    keyspaceTableList:[{keyspace:$ks, tableNameList:$tl, tableUUIDList:$tul}],
    timeBeforeDelete:0
  }')
CREATE=$(curl -sS -X POST "$YBA_URL/api/v1/customers/$CUUID/backups" -H "$AUTH" -H "$CT" -d "$BACKUP_PAYLOAD")
TASK_ID=$(echo "$CREATE" | jq -r '.taskUUID // .taskId // empty')
[[ -n "$TASK_ID" ]] || { echo "Backup create did not return taskUUID"; echo "$CREATE" | jq . || echo "$CREATE"; exit 1; }

echo "==> Waiting for backup task $TASK_ID"
wait_task "backup" "$TASK_ID" || { echo "Backup failed"; exit 2; }

# ---------- 2) Find this backup by taskUUID; extract defaultLocation
echo "==> Finding backup by taskUUID via /backups/page"
RAW=$(curl -sS -X POST "$YBA_URL/api/v1/customers/$CUUID/backups/page" \
  -H "$AUTH" -H "$CT" \
  -d "$(jq -n --arg u "$SRC_UUUID" \
        '{sortBy:"createTime", direction:"DESC",
          filter:{universeUUIDList:[$u], showHidden:true},
          limit:50, offset:0, needTotalCount:true}')")
ENTITY=$(echo "$RAW" | jq -r --arg tid "$TASK_ID" '.entities | map(select(.commonBackupInfo.taskUUID==$tid)) | first')
[[ "$ENTITY" != "null" && -n "$ENTITY" ]] || { echo "Could not locate backup entity for task $TASK_ID"; echo "$RAW" | jq .; exit 4; }
RESP_OBJ=$(echo "$ENTITY" | jq '(.commonBackupInfo.responseList | if type=="array" then .[0] else . end)')
LOC=$(echo "$RESP_OBJ" | jq -r '.defaultLocation // empty')
BID=$(echo "$ENTITY" | jq -r '.commonBackupInfo.backupUUID // empty')
[[ -n "$LOC" ]] || { echo "Could not extract defaultLocation"; echo "$ENTITY" | jq .; exit 5; }
echo "   location:   $LOC"
[[ -n "$BID" ]] && echo "   backupUUID: $BID"

# ---------- 2.5) Optional pre-drop ON TARGET (after backup success)
if $PRE_DROP; then
  [[ -n "$TARGET_YCQL_HOST" ]] || { echo "Need --target-ycql-host for --pre-drop"; exit 1; }
  command -v ycqlsh >/dev/null || { echo "ycqlsh required for --pre-drop" >&2; exit 1; }
  host="${TARGET_YCQL_HOST%%:*}"
  port="${TARGET_YCQL_HOST##*:}"; [[ "$host" == "$port" ]] && port="9042"
  YCQLSH_CMD=$(make_ycqlsh_cmd "$host" "$port")
  echo "==> Dropping target tables on $host:$port"
  for t in "${TABLES_ARR[@]}"; do
    echo "   DROP TABLE IF EXISTS ${KEYSPACE}.${t};"
    eval "$YCQLSH_CMD -e \"DROP TABLE IF EXISTS ${KEYSPACE}.${t};\"" || true
  done
fi

# ---------- 3) Restore only the requested tables (names + UUIDs)
RESTORE_PAYLOAD=$(jq -n \
  --arg action "RESTORE" --arg sc "$SCUUID" --arg tgt "$TGT_UUUID" --arg cu "$CUUID" \
  --arg ks "$KEYSPACE" --arg loc "$LOC" --argjson tl "$TABLES_JSON" --argjson tul "$TBL_UUIDS_JSON" \
  --argjson sse "$SSE" --argjson ut "$USE_TABLESPACES" '
  {actionType:$action, storageConfigUUID:$sc, universeUUID:$tgt,
   backupStorageInfoList:[{backupType:"YQL_TABLE_TYPE",keyspace:$ks,sse:$sse,storageLocation:$loc,
                           useTablespaces:$ut,tableNameList:$tl,tableUUIDList:$tul}],
   customerUUID:$cu}')
echo "==> Restoring selected tables via POST /api/v1/customers/$CUUID/restore"
RSP=$(curl -sS -w '\n%{http_code}' -X POST "$YBA_URL/api/v1/customers/$CUUID/restore" -H "$AUTH" -H "$CT" -d "$RESTORE_PAYLOAD")
BODY=$(echo "$RSP" | sed '$d'); CODE=$(echo "$RSP" | tail -n1)
[[ "$CODE" == "200" || "$CODE" == "202" ]] || { echo "Restore create failed (HTTP $CODE). Body:"; echo "$BODY" | jq . || echo "$BODY"; exit 6; }
RTASK=$(echo "$BODY" | jq -r '.taskUUID // .taskId // empty')
[[ -n "$RTASK" ]] || { echo "Restore response had no taskUUID"; echo "$BODY" | jq . || echo "$BODY"; exit 7; }
echo "==> Waiting for restore task $RTASK"
wait_task "restore" "$RTASK" || { echo "Restore failed"; exit 8; }

# ---------- 4) Optional delete (via POST /backups/delete)
if $DELETE_AFTER && [[ -n "$BID" ]]; then
  echo "==> Deleting backup $BID"
  DELETE_PAYLOAD=$(jq -n --arg bid "$BID" --arg sc "$SCUUID" \
    '{deleteBackupInfos:[{backupUUID:$bid, storageConfigUUID:$sc}]}')
  DEL=$(curl -sS -w '\n%{http_code}' -X POST "$YBA_URL/api/v1/customers/$CUUID/backups/delete" -H "$AUTH" -H "$CT" -d "$DELETE_PAYLOAD")
  DEL_BODY=$(echo "$DEL" | sed '$d'); DEL_CODE=$(echo "$DEL" | tail -n1)
  if [[ "$DEL_CODE" != "200" && "$DEL_CODE" != "202" ]]; then
    echo "Delete request failed (HTTP $DEL_CODE). Body:"; echo "$DEL_BODY" | jq . || echo "$DEL_BODY"
  else
    for dt in $(echo "$DEL_BODY" | jq -r '.taskUUID[]? // empty'); do
      echo "   Waiting for delete task $dt"
      wait_task "delete" "$dt" || { echo "Delete task $dt failed"; exit 9; }
    done
    if [[ -z "$(echo "$DEL_BODY" | jq -r '.taskUUID[]? // empty')" ]]; then
      maybe=$(echo "$DEL_BODY" | jq -r '.taskUUID // empty')
      if [[ -n "$maybe" && "$maybe" != "null" ]]; then
        echo "   Waiting for delete task $maybe"
        wait_task "delete" "$maybe" || { echo "Delete task $maybe failed"; exit 9; }
      fi
    fi
  fi
fi

echo "🎉 YCQL selective table refresh complete (subset backup + subset restore via UUIDs)."