#!/usr/bin/env bash
# Benchmark: (select auth.uid()) vs auth.uid() in Supabase-style RLS policies.
#
# Starts a throwaway container, loads setup.sql, then for each index state
# (no index / btree on user_id):
#   1. saves EXPLAIN (ANALYZE, BUFFERS, VERBOSE) for every variant x scenario
#   2. one warm-up pass over all combinations
#   3. RUNS interleaved timing passes (variant order rotated per run), \timing
#   4. counts function calls with a non-inlinable PL/pgSQL copy of auth.uid()
# Writes raw/<tag>/..., results-<tag>.csv. Removes the container at the end
# unless KEEP=1.
#
# Env: IMAGE (default postgres:17), NAME (default authuid-bench), RUNS (5),
#      VARIANTS (space-separated subset), SCENARIOS (subset of "A B C"),
#      INDEXES (subset of "noidx idx"), TAG (output label, default from IMAGE).
set -euo pipefail
cd "$(dirname "$0")"

IMAGE=${IMAGE:-postgres:17}
NAME=${NAME:-authuid-bench}
RUNS=${RUNS:-5}
SCENARIOS=${SCENARIOS:-"A B C"}
INDEXES=${INDEXES:-"noidx idx"}
TAG=${TAG:-pg${IMAGE#postgres:}}
ALL_VARIANTS="p1 p2 p3 p4 p4b own_p1 own_p2 own_p4 own_p4b own_lit"
VARIANTS=${VARIANTS:-$ALL_VARIANTS}

U=$(python3 -I -c "import hashlib,uuid;print(uuid.UUID(hashlib.md5(b'u42').hexdigest()))")
CLAIMS="{\"sub\":\"$U\",\"role\":\"authenticated\"}"
OUT=raw/$TAG
mkdir -p "$OUT/explain"

# Policy / WHERE expressions.
expr() {
  case $1 in
    p1|own_p1)   echo "user_id = auth.uid()" ;;
    p2|own_p2)   echo "user_id = (select auth.uid())" ;;
    p3)          echo "auth.uid() = user_id" ;;
    p4|own_p4)   echo "(select auth.jwt()) ->> 'role' = 'authenticated' and user_id = (select auth.uid())" ;;
    p4b|own_p4b) echo "auth.jwt() ->> 'role' = 'authenticated' and user_id = auth.uid()" ;;
    own_lit)     echo "user_id = '$U'" ;;
  esac
}

# Query text for a variant and scenario.
query() {
  local v=$1 s=$2 e
  e=$(expr "$v")
  if [[ $v == own_* ]]; then
    case $s in
      A) echo "select count(*) from todos where $e" ;;
      B) echo "select * from todos where done = false and ($e) order by id desc limit 20" ;;
      C) echo "select count(*) from todos where user_id = '$U' and ($e)" ;;
    esac
  else
    case $s in
      A) echo "select count(*) from todos" ;;
      B) echo "select * from todos where done = false order by id desc limit 20" ;;
      C) echo "select count(*) from todos where user_id = '$U'" ;;
    esac
  fi
}

# SQL that puts the session in the right state for a variant:
# RLS variants recreate the single policy and switch to role authenticated
# (the way PostgREST does: role + request.jwt.claims GUC);
# owner variants stay superuser/table owner (RLS not applied).
prelude() {
  local v=$1
  echo "reset role;"
  if [[ $v == own_* ]]; then
    echo "drop policy if exists p on todos;"
  else
    echo "drop policy if exists p on todos;"
    echo "create policy p on todos for select to authenticated using ($(expr "$v"));"
    echo "set role authenticated;"
  fi
}

psqlc() { docker exec -i "$NAME" psql -U postgres -X -q -v ON_ERROR_STOP=1 "$@"; }

if docker ps -a --format '{{.Names}}' | grep -qx "$NAME"; then
  echo "container $NAME exists; set NAME=... to use another name" >&2; exit 1
fi
docker run -d --name "$NAME" -e POSTGRES_PASSWORD=x "$IMAGE" >/dev/null
until docker exec "$NAME" pg_isready -U postgres -q 2>/dev/null; do sleep 1; done
sleep 2
until docker exec "$NAME" psql -U postgres -Atc 'select 1' >/dev/null 2>&1; do sleep 1; done
psqlc -At -c "select version()" | tee "$OUT/version.txt"
psqlc < setup.sql > "$OUT/setup.log" 2>&1

# Session start: same claims for the whole session (PostgREST uses
# set_config(..., true) inside the request transaction; session-level here).
SESSION="select set_config('request.jwt.claims', '$CLAIMS', false) \\gset x_"

: > "$OUT/timings.txt"
for IDX in $INDEXES; do
  if [[ $IDX == idx ]]; then
    psqlc -c "create index if not exists todos_user_id_idx on todos (user_id); analyze todos;"
  else
    psqlc -c "drop index if exists todos_user_id_idx; analyze todos;"
  fi

  # 1. EXPLAIN per combination
  for v in $VARIANTS; do for s in $SCENARIOS; do
    { echo "$SESSION"; prelude "$v"
      echo "explain (analyze, buffers, verbose) $(query "$v" "$s");"
    } | psqlc > "$OUT/explain/${IDX}_${v}_${s}.txt"
  done; done

  # 2+3. warm-up (run 0) and timed runs, one psql session per index state
  {
    echo "$SESSION"
    read -r -a VA <<< "$VARIANTS"
    n=${#VA[@]}
    for r in $(seq 0 "$RUNS"); do
      for k in $(seq 0 $((n-1))); do
        v=${VA[$(( (k + r) % n ))]}       # rotate order per run
        prelude "$v"
        for s in $SCENARIOS; do
          echo "\\echo M|$IDX|$r|$v|$s"
          echo "\\timing on"
          echo "$(query "$v" "$s");"
          echo "\\timing off"
        done
      done
    done
    echo "reset role; drop policy if exists p on todos;"
  } | psqlc -At | grep -E '^(M\||Time:)' >> "$OUT/timings.txt"

  # 4. call counts with the non-inlinable copy (track_functions = all)
  for pv in "user_id = auth.uid_pl()" "user_id = (select auth.uid_pl())"; do
    for s in $SCENARIOS; do
      q=$(query p1 "$s")
      psqlc -c "drop policy if exists p on todos; create policy p on todos for select to authenticated using ($pv);
                select pg_stat_reset_single_function_counters('auth.uid_pl'::regproc);
                select pg_stat_reset_single_function_counters('auth.uid'::regproc);" > /dev/null
      # own backend, flushes its function stats on exit
      { echo "set track_functions = 'all';"; echo "$SESSION"; echo "set role authenticated;"
        echo "$q;"; } | psqlc -At > /dev/null
      sleep 1
      calls=$(psqlc -At -c "select coalesce((select calls from pg_stat_user_functions where funcid = 'auth.uid_pl'::regproc), 0)")
      calls_uid=$(psqlc -At -c "select coalesce((select calls from pg_stat_user_functions where funcid = 'auth.uid'::regproc), 0)")
      echo "$IDX|$pv|$s|uid_pl_calls=$calls|uid_calls=$calls_uid" | tee -a "$OUT/callcounts.txt"
    done
  done
  # also: the real (inlined) auth.uid() under track_functions = all
  psqlc -c "drop policy if exists p on todos; create policy p on todos for select to authenticated using (user_id = auth.uid());
            select pg_stat_reset_single_function_counters('auth.uid'::regproc);" > /dev/null
  { echo "set track_functions = 'all';"; echo "$SESSION"; echo "set role authenticated;"
    echo "select count(*) from todos;"; } | psqlc -At > /dev/null
  sleep 1
  echo "$IDX|user_id = auth.uid() (real, sql)|A|uid_calls=$(psqlc -At -c "select coalesce((select calls from pg_stat_user_functions where funcid = 'auth.uid'::regproc), 0)")" | tee -a "$OUT/callcounts.txt"
  psqlc -c "drop policy if exists p on todos;"
done

# timings.txt -> csv
python3 -I - "$OUT/timings.txt" "results-$TAG.csv" "$(cat "$OUT/version.txt")" <<'PY'
import sys, csv
src, dst, ver = sys.argv[1:4]
rows, cur = [], None
for line in open(src):
    line = line.strip()
    if line.startswith("M|"):
        cur = line.split("|")[1:]
    elif line.startswith("Time:") and cur:
        ms = float(line.split()[1])
        idx, run, v, s = cur
        rows.append([ver.split(" on ")[0], idx, v, s, int(run), ms])
        cur = None
with open(dst, "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["version", "index", "variant", "scenario", "run", "ms"])
    w.writerows(rows)
print(f"{len(rows)} timings -> {dst}")
PY

[[ ${KEEP:-0} == 1 ]] || docker rm -f "$NAME" >/dev/null