Set up Thanos Query and Compactor for distributed metrics querying

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

# ---------------------------------------------------------------------------
# Thanos component installer: Sidecar, Query, Store Gateway, Compactor
# ---------------------------------------------------------------------------

VERSION="0.36.1"
THANOS_USER="thanos"
THANOS_GROUP="thanos"
INSTALL_DIR="/opt/thanos-${VERSION}.linux-amd64"
CONFIG_DIR="/etc/thanos"
CERTS_DIR="${CONFIG_DIR}/certs"
DATA_DIR="/var/lib/thanos"

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
log()  { echo -e "${GREEN}$*${NC}"; }
warn() { echo -e "${YELLOW}$*${NC}"; }
err()  { echo -e "${RED}$*${NC}" >&2; }

usage() {
    cat <<EOF
Usage: $0 --role <sidecar|query|store|compactor> [options]

Common options:
  --bucket-endpoint <host:port>   MinIO/S3 endpoint (required for all roles except query)
  --access-key <key>              S3 access key
  --secret-key <secret>           S3 secret key
  --bucket-name <name>            Bucket name (default: thanos-metrics)

Role-specific:
  sidecar:   --prometheus-url <url> --tsdb-path <path>
  query:     --store-addrs "host1:port,host2:port"  (comma separated gRPC endpoints)

Example:
  $0 --role sidecar --bucket-endpoint 203.0.113.20:9000 --access-key thanos-svc \\
     --secret-key 'secret' --prometheus-url http://127.0.0.1:9090 --tsdb-path /var/lib/prometheus
EOF
    exit 1
}

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------
ROLE=""
BUCKET_ENDPOINT=""
ACCESS_KEY=""
SECRET_KEY=""
BUCKET_NAME="thanos-metrics"
PROM_URL="http://127.0.0.1:9090"
TSDB_PATH="/var/lib/prometheus"
STORE_ADDRS=""

while [ $# -gt 0 ]; do
    case "$1" in
        --role) ROLE="$2"; shift 2 ;;
        --bucket-endpoint) BUCKET_ENDPOINT="$2"; shift 2 ;;
        --access-key) ACCESS_KEY="$2"; shift 2 ;;
        --secret-key) SECRET_KEY="$2"; shift 2 ;;
        --bucket-name) BUCKET_NAME="$2"; shift 2 ;;
        --prometheus-url) PROM_URL="$2"; shift 2 ;;
        --tsdb-path) TSDB_PATH="$2"; shift 2 ;;
        --store-addrs) STORE_ADDRS="$2"; shift 2 ;;
        -h|--help) usage ;;
        *) err "Unknown argument: $1"; usage ;;
    esac
done

[ -z "$ROLE" ] && { err "Missing --role"; usage; }
case "$ROLE" in
    sidecar|query|store|compactor) ;;
    *) err "Invalid role: $ROLE"; usage ;;
esac

if [ "$ROLE" != "query" ] && { [ -z "$BUCKET_ENDPOINT" ] || [ -z "$ACCESS_KEY" ] || [ -z "$SECRET_KEY" ]; }; then
    err "Roles other than 'query' require --bucket-endpoint, --access-key, --secret-key"
    usage
fi

if [ "$(id -u)" -ne 0 ]; then
    err "This script must be run as root or with sudo."
    exit 1
fi

# ---------------------------------------------------------------------------
# Distro detection
# ---------------------------------------------------------------------------
if [ -f /etc/os-release ]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian) PKG_MGR="apt"; PKG_INSTALL="apt install -y" ;;
        almalinux|rocky|centos|rhel|ol|fedora) PKG_MGR="dnf"; PKG_INSTALL="dnf install -y" ;;
        amzn) PKG_MGR="yum"; PKG_INSTALL="yum install -y" ;;
        *) err "Unsupported distro: $ID"; exit 1 ;;
    esac
else
    err "/etc/os-release not found - cannot detect distro"
    exit 1
fi

NOLOGIN_SHELL="/usr/sbin/nologin"
[ "$PKG_MGR" != "apt" ] && NOLOGIN_SHELL="/sbin/nologin"

TOTAL_STEPS=7
CURRENT=0
step() { CURRENT=$((CURRENT + 1)); echo -e "${GREEN}[${CURRENT}/${TOTAL_STEPS}]${NC} $*"; }

# ---------------------------------------------------------------------------
# Rollback on failure
# ---------------------------------------------------------------------------
SERVICE_NAME="thanos-${ROLE}"
cleanup_on_error() {
    err "Installation failed. Rolling back..."
    systemctl stop "${SERVICE_NAME}" 2>/dev/null || true
    systemctl disable "${SERVICE_NAME}" 2>/dev/null || true
    rm -f "/etc/systemd/system/${SERVICE_NAME}.service"
    systemctl daemon-reload 2>/dev/null || true
    exit 1
}
trap cleanup_on_error ERR

# ---------------------------------------------------------------------------
# 1. Prerequisites
# ---------------------------------------------------------------------------
step "Checking prerequisites and installing curl/tar..."
if [ "$PKG_MGR" = "apt" ]; then
    apt update
fi
$PKG_INSTALL curl tar >/dev/null

# ---------------------------------------------------------------------------
# 2. Create system user
# ---------------------------------------------------------------------------
step "Creating dedicated thanos system user..."
if ! id "$THANOS_USER" >/dev/null 2>&1; then
    useradd --system --no-create-home --shell "$NOLOGIN_SHELL" "$THANOS_USER"
    log "Created user ${THANOS_USER}"
else
    warn "User ${THANOS_USER} already exists, skipping"
fi

# ---------------------------------------------------------------------------
# 3. Download and install Thanos binary
# ---------------------------------------------------------------------------
step "Downloading and installing Thanos v${VERSION}..."
if [ ! -x /usr/local/bin/thanos ]; then
    TMP_TARBALL="$(mktemp /tmp/thanos.XXXXXX.tar.gz)"
    curl -fL -o "$TMP_TARBALL" \
        "https://github.com/thanos-io/thanos/releases/download/v${VERSION}/thanos-${VERSION}.linux-amd64.tar.gz"
    tar -xzf "$TMP_TARBALL" -C /opt
    rm -f "$TMP_TARBALL"
    ln -sf "${INSTALL_DIR}/thanos" /usr/local/bin/thanos
else
    warn "thanos binary already installed, skipping download"
fi

# ---------------------------------------------------------------------------
# 4. Config, cert, and data directories with correct ownership/permissions
# ---------------------------------------------------------------------------
step "Setting up config, cert, and data directories..."
mkdir -p "$CERTS_DIR" "$DATA_DIR"
chown -R "${THANOS_USER}:${THANOS_GROUP}" "$CONFIG_DIR" "$DATA_DIR"
chmod 755 "$CONFIG_DIR" "$CERTS_DIR" "$DATA_DIR"

# Object storage config file - only needed for sidecar/store/compactor
if [ "$ROLE" != "query" ]; then
    BUCKET_FILE="${CONFIG_DIR}/bucket.yml"
    cat > "$BUCKET_FILE" <<EOF
type: S3
config:
  bucket: "${BUCKET_NAME}"
  endpoint: "${BUCKET_ENDPOINT}"
  access_key: "${ACCESS_KEY}"
  secret_key: "${SECRET_KEY}"
  insecure: false
  signature_version2: false
  http_config:
    tls_config:
      ca_file: ${CERTS_DIR}/minio-ca.pem
EOF
    chown "${THANOS_USER}:${THANOS_GROUP}" "$BUCKET_FILE"
    chmod 640 "$BUCKET_FILE"
fi

# ---------------------------------------------------------------------------
# 5. Build the ExecStart line per role
# ---------------------------------------------------------------------------
step "Generating systemd unit for thanos-${ROLE}..."
case "$ROLE" in
    sidecar)
        EXEC_LINE="/usr/local/bin/thanos sidecar \\
  --tsdb.path=${TSDB_PATH} \\
  --prometheus.url=${PROM_URL} \\
  --grpc-address=0.0.0.0:10901 \\
  --http-address=127.0.0.1:10902 \\
  --objstore.config-file=${CONFIG_DIR}/bucket.yml \\
  --grpc-server-tls-cert=${CERTS_DIR}/sidecar.pem \\
  --grpc-server-tls-key=${CERTS_DIR}/sidecar-key.pem \\
  --grpc-server-tls-client-ca=${CERTS_DIR}/ca.pem"
        # Sidecar needs to read Prometheus data directory
        usermod -aG prometheus "$THANOS_USER" 2>/dev/null || warn "prometheus group not found, skipping group add"
        ;;
    query)
        [ -z "$STORE_ADDRS" ] && { err "--store-addrs required for query role"; usage; }
        IFS=',' read -ra ADDR_ARR <<< "$STORE_ADDRS"
        STORE_FLAGS=""
        for a in "${ADDR_ARR[@]}"; do
            STORE_FLAGS+=" --store=${a}"
        done
        EXEC_LINE="/usr/local/bin/thanos query \\
  --grpc-address=0.0.0.0:10901 \\
  --http-address=0.0.0.0:10904 \\
  --query.replica-label=replica${STORE_FLAGS} \\
  --grpc-client-tls-cert=${CERTS_DIR}/query.pem \\
  --grpc-client-tls-key=${CERTS_DIR}/query-key.pem \\
  --grpc-client-tls-ca=${CERTS_DIR}/ca.pem"
        ;;
    store)
        EXEC_LINE="/usr/local/bin/thanos store \\
  --data-dir=${DATA_DIR}/store \\
  --grpc-address=0.0.0.0:10901 \\
  --http-address=127.0.0.1:10905 \\
  --objstore.config-