testbedctl

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

PROJ_DIR="$(cd "$(dirname "$0")" && pwd)"

usage() {
  cat <<EOF
Usage: testbedctl <command> [args...]

Commands:
  psql          Connect to GreptimeDB via PostgreSQL protocol
  mysql         Connect to GreptimeDB via MySQL protocol
  s3  [args..]  Run aws s3 against Garage (auto-sources credentials)
  telemetrygen up              Start OTLP trace ingestion
  telemetrygen down            Stop trace ingestion
  telemetrygen metrics up      Start continuous OTLP metrics (~50 randomized names, Python OTel)
  telemetrygen metrics down    Stop metrics ingestion
  metrics-partition  Create greptime_physical_table partitioned on 'timebox' (4 partitions)
  duckdb        Open DuckDB shell with Iceberg REST catalog attached
  pyiceberg [table] [-local]
                Query Iceberg tables via pyiceberg.
                Pass a table name (default: opentelemetry_traces4).
                Use -local for standalone-fs (local filesystem, no S3).
  flush <table> Flush a table's memtable (admin flush_table) to trigger manifest export
  clean         Remove .greptimedb directory
  help          Show this help
EOF
}

source_s3_env() {
  local s3_env="${PROJ_DIR}/.greptimedb/s3.env"
  if [ -f "$s3_env" ]; then
    # shellcheck disable=SC1090
    . "$s3_env"
  else
    echo "WARNING: ${s3_env} not found. Run the cluster first." >&2
  fi
}

cmd_psql() {
  # Client PG port 11043: haproxy (distributed) or standalone/standalone-fs.
  exec psql -h 127.0.0.1 -p 11043 -d public -U root "$@"
}

cmd_flush() {
  local table="${1:-}"
  if [ -z "$table" ]; then
    echo "Usage: testbedctl flush <table_name>" >&2
    exit 1
  fi
  # Client PG port 11043: haproxy (distributed) or standalone/standalone-fs.
  exec psql -h 127.0.0.1 -p 11043 -d public -U root \
    -c "admin flush_table('${table}');"
}

cmd_mysql() {
  # Client MySQL port 11042: haproxy (distributed) or standalone/standalone-fs.
  exec mysql -h 127.0.0.1 -P 11042 -u root "$@"
}

cmd_s3() {
  source_s3_env
  exec aws s3 "$@"
}

TG_DIR="${PROJ_DIR}/datasources/telemetrygen"
TG_COMPOSE="${TG_DIR}/compose.yaml"        # traces service only
TG_METRICS_GEN="${TG_DIR}/gen_metrics.py"   # Python OTLP metrics generator (single process)
TG_METRICS_PID="${PROJ_DIR}/.greptimedb/tg-metrics.pid"
TG_METRICS_LOG="${PROJ_DIR}/.greptimedb/tg-metrics.log"

tg_metrics_running() {
  [ -f "$TG_METRICS_PID" ] && kill -0 "$(cat "$TG_METRICS_PID")" 2>/dev/null
}

cmd_telemetrygen_metrics() {
  local action="${1:-}"
  case "$action" in
    up)
      mkdir -p "${PROJ_DIR}/.greptimedb"
      if tg_metrics_running; then
        echo "==> telemetrygen metrics already running (pid $(cat "$TG_METRICS_PID"))"
        echo "    log: $TG_METRICS_LOG"
        return
      fi
      rm -f "$TG_METRICS_PID"
      echo "==> starting telemetrygen metrics" \
           "(count=${TG_METRIC_COUNT:-50}, seed=${TG_METRIC_SEED:-0}," \
           "workers=${TG_METRICS_WORKERS:-2}, rate=${TG_METRICS_RATE:-10}/s/worker)"
      # Single background Python OTLP process; inherits all TG_* env vars.
      nohup python3 "$TG_METRICS_GEN" >> "$TG_METRICS_LOG" 2>&1 &
      local pid=$!
      echo "$pid" > "$TG_METRICS_PID"
      # Brief health check: if it died immediately (e.g. missing opentelemetry
      # deps because not run under `nix develop`), surface the log.
      sleep 2
      if ! kill -0 "$pid" 2>/dev/null; then
        echo "ERROR: telemetrygen metrics exited immediately. Recent log:" >&2
        tail -n 20 "$TG_METRICS_LOG" >&2
        rm -f "$TG_METRICS_PID"
        echo >&2 "Hint: run inside 'nix develop' so python3 has opentelemetry-sdk." >&2
        return 1
      fi
      echo "==> running (pid $pid)  ->  ${TG_OTLP_ENDPOINT:-127.0.0.1:11040}"
      echo "    log: $TG_METRICS_LOG"
      echo "==> metric names (name <TAB> type):"
      python3 "$TG_METRICS_GEN" --list | sed 's/^/    /'
      ;;
    down)
      if [ ! -f "$TG_METRICS_PID" ]; then
        echo "==> telemetrygen metrics not running (no pidfile)"
        return
      fi
      local pid
      pid="$(cat "$TG_METRICS_PID")"
      if kill -0 "$pid" 2>/dev/null; then
        kill -TERM "$pid" 2>/dev/null || true
        # Give it a moment to flush a final export on graceful shutdown.
        local i
        for i in $(seq 1 20); do
          kill -0 "$pid" 2>/dev/null || break
          sleep 0.25
        done
        if kill -0 "$pid" 2>/dev/null; then
          echo "==> process did not exit; sending SIGKILL"
          kill -KILL "$pid" 2>/dev/null || true
        fi
        echo "==> stopped telemetrygen metrics (pid $pid)"
      else
        echo "==> pid $pid not alive (stale pidfile)"
      fi
      rm -f "$TG_METRICS_PID"
      ;;
    *)
      cat >&2 <<EOF
Usage: testbedctl telemetrygen metrics <up|down>

Runs a SINGLE Python OTLP process that emits '${TG_METRIC_COUNT:-50}' distinct
metrics, each with a randomized name and a randomized type (Gauge/Sum/Histogram)
— replacing the old many-container telemetrygen setup.

Tunables (env):
  TG_METRIC_COUNT     number of distinct metric names (default 50)
  TG_METRIC_SEED      RNG seed for reproducible names (default 0; change to reshuffle)
  TG_OTLP_ENDPOINT    OTLP host:port (default 127.0.0.1:11040 = haproxy in distributed, standalone in single-node)
  TG_OTLP_URL_PATH    OTLP HTTP path (default /v1/otlp/v1/metrics)
  TG_METRICS_INTERVAL export interval (default 5s)
  TG_METRICS_RATE     observations/sec per worker (default 10)
  TG_METRICS_WORKERS  worker threads (default 2)
EOF
      [ -n "$action" ] && echo "Unknown action: $action" >&2
      exit 1
      ;;
  esac
}

cmd_telemetrygen() {
  case "${1:-}" in
    metrics)
      shift
      cmd_telemetrygen_metrics "$@"
      ;;
    up)
      # Detached: returns immediately; the finite trace burst runs to completion.
      exec podman-compose -f "$TG_COMPOSE" up -d telemetrygen
      ;;
    down)
      exec podman rm -f --ignore telemetrygen
      ;;
    *)
      cat >&2 <<EOF
Usage:
  testbedctl telemetrygen up            Start OTLP trace ingestion
  testbedctl telemetrygen down          Stop trace ingestion
  testbedctl telemetrygen metrics up    Start continuous metrics (~50 randomized metric names)
  testbedctl telemetrygen metrics down  Stop metrics ingestion

Metrics tunables (env): TG_METRIC_COUNT, TG_METRIC_SEED, TG_OTLP_ENDPOINT,
                       TG_METRICS_INTERVAL, TG_METRICS_RATE, TG_METRICS_WORKERS
EOF
      [ -n "${1:-}" ] && echo "Unknown telemetrygen action: $1" >&2
      exit 1
      ;;
  esac
}

cmd_metrics_partition() {
  # Sets up the metric-engine physical table with 4 range partitions on the
  # 'timebox' column. telemetrygen emits 'timebox' via its --unique-timeseries
  # flag (enabled in the metrics compose), giving a semi-random partition key.
  local port="${TG_DB_PORT:-11043}"
  local PSQL=(psql -h 127.0.0.1 -p "$port" -d public -U root -t -A -v ON_ERROR_STOP=1)

  phys_exists() {
    [ "$("${PSQL[@]}" -c "SELECT 1 FROM information_schema.tables WHERE table_name='greptime_physical_table'" 2>/dev/null)" = "1" ]
  }

  if phys_exists; then
    echo "==> greptime_physical_table exists; attempting to drop it (its metric data will be lost)..."
    local drop_out
    drop_out="$("${PSQL[@]}" -c "DROP TABLE greptime_physical_table" 2>&1)" || true
    if phys_exists; then
      echo "$drop_out" >&2
      cat >&2 <<'EOF'
Cannot replace 'greptime_physical_table': it still backs live logical metric tables
(e.g. telemetrygen_* metrics). Reset to a clean database first:

    process-compose down && ./testbedctl clean && process-compose up haproxy     # distributed cluster
    # or:  process-compose up standalone-fs                                   # single node

then re-run:  ./testbedctl metrics-partition
EOF
      exit 1
    fi
    echo "==> dropped existing greptime_physical_table"
  fi

  echo "==> creating greptime_physical_table with 4 partitions on the 'timebox' column..."
  "${PSQL[@]}" 2>&1 <<'SQL'
CREATE TABLE greptime_physical_table (
  greptime_timestamp TIMESTAMP NOT NULL,
  greptime_value DOUBLE NULL,
  timebox STRING NULL,
  TIME INDEX (greptime_timestamp),
  PRIMARY KEY (timebox)
)
PARTITION ON COLUMNS (timebox) (
  timebox < '2',
  timebox >= '2' AND timebox < '5',
  timebox >= '5' AND timebox < '8',
  timebox >= '8'
)
ENGINE = metric
WITH ('physical_metric_table' = 'true');
SQL

  if phys_exists; then
    echo "==> done. greptime_physical_table now has 4 partitions keyed on 'timebox'."
    echo "==> populate it with:"
    echo "      ./testbedctl telemetrygen metrics up"
    echo "    (the metrics compose uses --unique-timeseries, which emits 'timebox')"
    echo "==> note: once partitioned, only metrics carrying 'timebox' route correctly;"
    echo "    avoid ingesting other metric sources on the same database (they lack 'timebox')."
  else
    echo "ERROR: greptime_physical_table was not created" >&2
    exit 1
  fi
}

cmd_duckdb() {
  source_s3_env
  local init_file
  init_file=$(mktemp)
  cat > "$init_file" <<DUCKINIT
SET autoinstall_known_extensions=1;
SET autoload_known_extensions=1;
FORCE INSTALL iceberg;
LOAD iceberg;
INSTALL httpfs;
LOAD httpfs;
SET s3_endpoint='127.0.0.1:11010';
SET s3_access_key_id='${S3_ACCESS_KEY}';
SET s3_secret_access_key='${S3_SECRET_KEY}';
SET s3_region='garage';
SET s3_use_ssl=false;
SET s3_url_style='path';

-- Iceberg REST catalog (SHOW ALL TABLES works)
ATTACH 'greptime' AS iceberg_cat (
    TYPE iceberg,
    ENDPOINT 'http://127.0.0.1:11040/v1/iceberg',
    AUTHORIZATION_TYPE 'none'
);

-- Workaround: query via read_parquet (DuckDB REST catalog client
-- has a bug with manifest path resolution in v1.5.2)
CREATE OR REPLACE VIEW otel_traces AS
  SELECT * FROM read_parquet(
    's3://test-bucket/data/data/greptime/public/*/*/*.parquet',
    union_by_name=true
  );
DUCKINIT
  echo "==> Iceberg catalog 'greptime' attached as 'iceberg_cat'"
  echo "==> Try:  SHOW ALL TABLES;"
  echo "==> Try:  SELECT * FROM otel_traces LIMIT 10;"
  echo "==> Try:  SELECT count(*) FROM otel_traces;"
  echo
  exec duckdb -init "$init_file" "$@"
}

cmd_pyiceberg() {
  # Scan all args for the -local flag so the table name can appear in any
  # position (e.g. `pyiceberg my_table -local` or `pyiceberg -local my_table`).
  local local_mode=false
  for a in "$@"; do
    if [ "$a" = "-local" ] || [ "$a" = "--local" ]; then
      local_mode=true
      break
    fi
  done
  if [ "$local_mode" = true ]; then
    echo "==> pyiceberg LOCAL mode (standalone-fs, no S3)"
  else
    source_s3_env
  fi
  exec python3 "${PROJ_DIR}/scripts/read_iceberg.py" "$@"
}

cmd_clean() {
  if [ -d "${PROJ_DIR}/.greptimedb" ]; then
    rm -rf "${PROJ_DIR}/.greptimedb"
    echo "==> Removed .greptimedb"
  else
    echo "==> .greptimedb not found, nothing to clean."
  fi
}

case "${1:-help}" in
  psql)         shift; cmd_psql "$@" ;;
  mysql)        shift; cmd_mysql "$@" ;;
  s3)           shift; cmd_s3 "$@" ;;
  telemetrygen) shift; cmd_telemetrygen "$@" ;;
  metrics-partition) shift; cmd_metrics_partition "$@" ;;
  duckdb)       shift; cmd_duckdb "$@" ;;
  pyiceberg)    shift; cmd_pyiceberg "$@" ;;
  flush)        shift; cmd_flush "$@" ;;
  clean)        cmd_clean ;;
  help|--help|-h)
    usage
    ;;
  *)
    echo "Unknown command: $1" >&2
    usage >&2
    exit 1
    ;;
esac