Script to snapshot a Prometheus pod and download the snapshot locally

prom_snapshot.sh

#!/bin/bash

# Ensure that the script exits on any errors
set -euo pipefail

if [[ -n "${VERBOSE:-}" ]]; then
    set -x
fi

# Check that the correct number of arguments are provided
if [ "$#" -ne 2 ]; then
    echo "Usage: $0 <prometheus_pod> <target_directory> [container name] [port]"
    exit 1
fi

# Function to cleanup port-forwarding on exit
cleanup() {
    if [[ -n "${PORT_FORWARD_PID:-}" ]]; then
        echo "Cleaning up..."
        kill "${PORT_FORWARD_PID}"
    fi
}
trap cleanup EXIT

PROMETHEUS_POD="${1}"
TARGET_DIR="${2}"
PROMETHEUS_CONTAINER="${3:-prometheus-server}"
REMOTE_PORT="${4:-9090}"
LOCAL_PORT=9090

PROMETHEUS_STORAGE_DIR="/prometheus"
# Sniff the Pod JSON to determine if `--storage.tsdb.path` is set and override the default storage path set above.
PROMETHEUS_STORAGE_DIR_OVERRIDE=$(kubectl get pod "${PROMETHEUS_POD}" -o json | jq -r --arg container "${PROMETHEUS_CONTAINER}" '.spec.containers[] | select(.name == $container) | .args[] | select(contains("storage.tsdb.path=")) | split("=")[1]')
if [[ -n "${PROMETHEUS_STORAGE_DIR_OVERRIDE}" ]]; then
    echo "Overriding storage path to ${PROMETHEUS_STORAGE_DIR_OVERRIDE}"
    PROMETHEUS_STORAGE_DIR="${PROMETHEUS_STORAGE_DIR_OVERRIDE}"
fi

# Remove any leading / from PROMETHUS_STORAGE_DIR
PROMETHEUS_STORAGE_DIR="${PROMETHEUS_STORAGE_DIR#/}"

# Start port-forwarding to the Prometheus pod
echo "Starting port-forward to ${PROMETHEUS_POD}..."
kubectl port-forward "${PROMETHEUS_POD}" "${LOCAL_PORT}:${REMOTE_PORT}" &
PORT_FORWARD_PID=$!

# Wait a few seconds for the port-forward to establish
sleep 5

# Trigger the snapshot via the Prometheus admin API
echo "Triggering snapshot..."
SNAPSHOT_RESPONSE=$(curl -s -XPOST "http://localhost:${LOCAL_PORT}/api/v1/admin/tsdb/snapshot")
SNAPSHOT_ERROR=$(echo "${SNAPSHOT_RESPONSE}" | jq -r '.error')
if [[ "${SNAPSHOT_ERROR}" != "null" ]]; then
    echo "Failed to create snapshot: ${SNAPSHOT_ERROR}"
    exit 1
fi

SNAPSHOT_NAME=$(echo -n "${SNAPSHOT_RESPONSE}" | jq -r '.data.name')
if [[ "${SNAPSHOT_NAME}" == "null" ]]; then
    echo "Failed to create snapshot"
    exit 1
fi

echo "Snapshot created: ${SNAPSHOT_NAME}"

# Copy the snapshot from the pod to the local machine
echo "Copying snapshot from pod to local directory..."

# Use kubectl cp to copy the snapshot directory to the target directory
# Keep retrying forever as this can get quite big!
kubectl cp --container "${PROMETHEUS_CONTAINER}" --retries=-1 "${PROMETHEUS_POD}:/${PROMETHEUS_STORAGE_DIR}/snapshots/${SNAPSHOT_NAME}" "${TARGET_DIR}"

echo "Snapshot copied to ${TARGET_DIR}"

# Cleanup will happen automatically via the trap