Cluster-Wide Catalog Table Query Aggregation

				
					#!/bin/bash

# Script:      yb_cluster_query.sh
# Description: This script will execute provided query on all nodes in a cluster and
#              dump the output to a file.
# Created:     04/11/2025
# Last update: 04/11/2025

# Call Spec:
#   ./yb_cluster_query.sh <<TServer IP>> <<Database Name>> <<Database Username>> <<SQL File>> <<Output File>> <<Delimeter>>

# Example Use:
#   ./yb_cluster_query.sh 127.0.0.1 yugabyte yugabyte pg_stat_statements.sql cluster_output.csv "|"

# Important Notes:
#  1. This script is for YSQL only
#  2. ysqlsh must be in the PATH
#  3. The SQL File must contain only one query
#  4. The delimiter is optional, and will default to "|"

# Expected input parameters
MASTER_HOST=$1
DATABASE=$2
USER=$3
SQL_FILE=$4
OUTPUT_FILE=$5
DELIM=${6:-"|"}

# Setup
PORT=5433

if [[ $# -lt 4 ]]; then
  echo "Usage: $0 <master_host> <port> <sql_file> <output_file> [delimiter]"
  exit 1
fi

# Get the query from the SQL file, strip the trailing semicolon if present
QUERY=$(<"$SQL_FILE")
QUERY=${QUERY%;}

# Get host list using ysqlsh in tuples-only and aligned-off mode
HOSTS=$(ysqlsh -h "$MASTER_HOST" -p "$PORT" -U "$USER" -d "$DATABASE" -Atc "SELECT host FROM yb_servers();")

# Create a temp dir for output
TMP_DIR=$(mktemp -d)

# Get header using LIMIT 0
HEADER_LINE=$(ysqlsh -h "$MASTER_HOST" -p "$PORT" -U "$USER" -d "$DATABASE" -F "$DELIM" -Ac "SELECT * FROM (${QUERY}) AS subq LIMIT 0;" 2>/dev/null | head -n 1)
echo "host${DELIM}${HEADER_LINE}" > "$TMP_DIR/header"

# Function to run query on a single host
run_query() {
  local HOST=$1
  local OUT_FILE="$TMP_DIR/${HOST}.out"
  ysqlsh -h "$HOST" -p "$PORT" -U "$USER" -d "$DATABASE" -AtF "$DELIM" -c "$QUERY" 2>/dev/null | sed "s/^/${HOST}${DELIM}/" > "$OUT_FILE"
  if [[ ! -s "$OUT_FILE" ]]; then
    echo "Warning: No output for $HOST" >&2
  fi
}

export -f run_query
export QUERY PORT DELIM TMP_DIR

# Run all queries in parallel
echo "$HOSTS" | xargs -n1 -P 10 -I {} bash -c 'run_query "$@"' _ {}

# Write header
cat "$TMP_DIR/header" > "$OUTPUT_FILE"

# Sort data rows and append
find "$TMP_DIR" -name '*.out' -exec cat {} + | sort >> "$OUTPUT_FILE"

# Cleanup
rm -rf "$TMP_DIR"

echo -e "\n✅ Results written to $OUTPUT_FILE"