#!/usr/bin/env bash
# scripts/prod-psql — psql into the LIVE precis_prod DB via a cluster node.
#
# scripts/db is LOCAL only (precis / precis_test). Prod sits behind pgbouncer
# on the cluster; reach it by hopping through a node (caspar/melchior) as
# agent_rw. Wraps the ssh-agent workaround (-o IdentityAgent=none) and the
# fixed pgbouncer coordinates so you stop retyping them.
#
# ⚠ agent_rw is WRITE-capable and this is PRODUCTION. Prefer read-only SELECTs.
#
# Usage:
#   scripts/prod-psql "SELECT count(*) FROM refs;"    # one-shot query
#   scripts/prod-psql "\d refs"                        # psql backslash cmds too
#   echo "SELECT ..." | scripts/prod-psql              # piped SQL (stdin)
#   scripts/prod-psql                                  # interactive shell
#
#   scripts/prod-psql --ro "SELECT ..."               # read-only: agent_ro role,
#                                                      # wrapped in BEGIN READ ONLY
#
# Row cap (.claude/fleet/authority.md, 2026-10-15): a one-shot or piped SQL
# text without --ro that mentions UPDATE or DELETE runs inside BEGIN; the
# rows the transaction updated or deleted (all its statements, see
# wrap_cap) are checked against PRECIS_PROD_WRITE_MAX_ROWS (default 100).
# Over the cap it ROLLBACKs and exits 4; under it, COMMIT. The count is
# printed either way. The match is on the raw text, so a read that merely
# mentions the word runs in the wrapper too (0 rows, commits). SQL that ends the
# transaction itself (BEGIN/COMMIT/ROLLBACK/END/ABORT) is refused, and a
# server-side marker catches one that slips past the refusal (exit 5).
#
# ⚠ Never `SET` a session setting here (`SET default_transaction_read_only`,
# `SET SESSION …`, `RESET …`). pgbouncer runs transaction pooling: a
# session-level SET stays on the SHARED server connection and every later
# client inherits it — on 2026-10-02 one such "read-only guard" made every
# agent_rw write fleet-wide fail until pgbouncer was reconnected. This script
# refuses those statements. For a read-only guard use --ro; inside your own
# transaction, `SET LOCAL` is fine (scripts/guide-web's comment says the same).
#
# Overrides (env):
#   PRECIS_PROD_SSH_HOST=melchior     # default caspar
#   PRECIS_PROD_PSQL_OPTS="-At"       # extra psql flags (e.g. terse -At, -x)
#   PRECIS_PROD_WRITE_MAX_ROWS=500    # row cap for UPDATE/DELETE (default 100)
set -euo pipefail

RO=0
if [[ "${1:-}" == "--ro" ]]; then
    RO=1
    shift
fi

HOST="${PRECIS_PROD_SSH_HOST:-caspar}"
# Resolved from the gitignored cluster overlay — the address must not live in
# this public repo (see scripts/lib/pgb-host.sh).
. "$(dirname "$0")/lib/pgb-host.sh"
PGB_HOST="$(resolve_pgb_host)"
PGB_PORT=6432
PGUSER=agent_rw
# --ro: the grants-based read-only role (SELECT-only; enforced per statement
# server-side, so pooling-proof — the same choice scripts/guide-web makes).
if [[ $RO -eq 1 ]]; then
    PGUSER=agent_ro
fi
PGDB=precis_prod
OPTS="${PRECIS_PROD_PSQL_OPTS:-}"

# Fixed connection string — no user input interpolated, so it's quoting-safe.
# -P pager=off and ON_ERROR_STOP make one-shot output clean and fail loudly.
REMOTE="psql -h ${PGB_HOST} -p ${PGB_PORT} -U ${PGUSER} -d ${PGDB} -P pager=off -v ON_ERROR_STOP=1 ${OPTS}"
# -X (skip the remote user's ~/.psqlrc) on the SCRIPTED paths only. That file
# is written for humans — on caspar it sets `\timing on`, `\x auto` and
# unicode `\pset` borders — and every one of those lands on STDOUT, so a
# caller parsing a `-At` scalar gets "1755892345\nTime: 0.4 ms" instead of a
# number. Bit the deploy drain (never once ran, 2026-08-22) and would bite
# scripts/deploy's canary heartbeat check on its first use. The interactive
# branch below deliberately keeps ~/.psqlrc — that's a human at a prompt.
REMOTE_SCRIPTED="psql -X -h ${PGB_HOST} -p ${PGB_PORT} -U ${PGUSER} -d ${PGDB} -P pager=off -v ON_ERROR_STOP=1 ${OPTS}"

# Refuse a session-level SET / RESET (see the ⚠ above). Statement-level, so
# `UPDATE … SET col = …` passes; `SET LOCAL`, `SET TRANSACTION` and
# `SET CONSTRAINTS` are transaction-scoped and pass too.
refuse_session_set() {
    if printf '%s' "$1" | perl -0777 -ne '
        s{--[^\n]*}{}g; s{/\*.*?\*/}{}gs;
        for my $st (split /;/) {
            next unless $st =~ /^\s*(SET|RESET)\b/i;
            next if $st =~ /^\s*SET\s+(LOCAL|TRANSACTION|CONSTRAINTS)\b/i;
            exit 1;
        }
        exit 0;'; then
        return 0
    fi
    echo "prod-psql: refused — a session-level SET/RESET through pgbouncer (transaction pooling) sticks to the shared server connection and poisons every later client. Use --ro for a read-only guard, or SET LOCAL inside BEGIN … COMMIT." >&2
    exit 2
}

# --ro wraps the input in one read-only transaction (belt to agent_ro's
# braces). A COMMIT inside the input would end it early; that is the
# caller's business, the role still refuses writes. An unterminated user
# statement would swallow the COMMIT ("syntax error at or near COMMIT"), so
# a missing trailing `;` is appended on its own line (safe after a trailing
# `-- comment`, and an empty statement is a no-op in psql).
wrap_ro() {
    if [[ $RO -eq 1 ]]; then
        local body="$1" trimmed
        trimmed="${body%"${body##*[![:space:]]}"}"
        if [[ "$trimmed" != *\; ]]; then
            body+=$'\n;'
        fi
        printf 'BEGIN READ ONLY;\n%s\nCOMMIT;\n' "$body"
    else
        printf '%s\n' "$1"
    fi
}

MAX_ROWS="${PRECIS_PROD_WRITE_MAX_ROWS:-100}"
if [[ ! "$MAX_ROWS" =~ ^[0-9]+$ ]]; then
    echo "prod-psql: PRECIS_PROD_WRITE_MAX_ROWS must be a whole number, got '$MAX_ROWS'." >&2
    exit 2
fi

# True when a non --ro SQL text mentions UPDATE or DELETE as a word (raw,
# any case): a superset of the statements that change existing rows.
needs_cap() {
    [[ $RO -eq 0 ]] &&
        printf '%s' "$1" | grep -qiE '(^|[^[:alnum:]_])(UPDATE|DELETE)([^[:alnum:]_]|$)'
}

# A capped write leaves the transaction to the wrapper, or its own COMMIT
# lands before the check.
refuse_tx_control() {
    if printf '%s' "$1" | perl -0777 -ne '
        s{--[^\n]*}{}g; s{/\*.*?\*/}{}gs;
        for my $st (split /;/) {
            exit 1 if $st =~ /^\s*(BEGIN|COMMIT|ROLLBACK|END|ABORT|START\s+TRANSACTION|PREPARE\s+TRANSACTION)\b/i;
        }
        exit 0;'; then
        return 0
    fi
    echo "prod-psql: refused — an UPDATE/DELETE runs inside the row-cap transaction; drop your own BEGIN/COMMIT/ROLLBACK." >&2
    exit 2
}

# BEGIN, a SET LOCAL marker, the user's SQL, then psql compares the rows
# the transaction updated or deleted to the cap and picks COMMIT or
# ROLLBACK. The count is the server's own per-transaction tally
# (pg_stat_xact_user_tables, readable by any role), not psql's ROW_COUNT:
# ROW_COUNT is the LAST query's, and a trailing `;;` or `; -- note` sends an
# empty query that resets it to 0 (measured 2026-10-15, psql 18 / PG 17).
# Cascaded and trigger-made changes count too. That view also carries counts
# a pooled server connection has not flushed yet from EARLIER clients'
# transactions (prod via pgbouncer, 2026-10-15: a 1-row UPDATE read 4, a
# 150-row one 152), so the cap is on the difference from a baseline read
# right after BEGIN; the server flushes only between transactions, so the
# baseline holds for ours. The marker is gone when the
# user's SQL ended the transaction early: nothing is left to roll back, and
# the wrapper says so. The "prod-psql: <VERDICT>" lines are what run_sql
# reads back. The trailing `;` rule is wrap_ro's.
wrap_cap() {
    local body="$1" trimmed
    trimmed="${body%"${body##*[![:space:]]}"}"
    if [[ "$trimmed" != *\; ]]; then
        body+=$'\n;'
    fi
    local tally='(SELECT coalesce(sum(n_tup_upd + n_tup_del), 0) AS n FROM pg_stat_xact_user_tables) s'
    printf 'BEGIN;\nSET LOCAL precis.prod_psql_cap = on;\n'
    printf 'SELECT n AS prod_psql_base FROM %s \\gset\n%s\n' "$tally" "$body"
    cat <<EOF
SELECT n - :prod_psql_base AS prod_psql_rows, n - :prod_psql_base > ${MAX_ROWS} AS prod_psql_over, coalesce(current_setting('precis.prod_psql_cap', true), '') = 'on' AS prod_psql_in_tx FROM ${tally} \\gset
\\if :prod_psql_in_tx
\\if :prod_psql_over
\\echo prod-psql: ROLLBACK, :prod_psql_rows rows is over the cap of ${MAX_ROWS} (PRECIS_PROD_WRITE_MAX_ROWS), nothing changed
ROLLBACK;
\\else
\\echo prod-psql: COMMIT, :prod_psql_rows rows, within the cap of ${MAX_ROWS}
COMMIT;
\\endif
\\else
\\echo prod-psql: UNCAPPED, the SQL ended the transaction itself, its writes are already committed
\\endif
EOF
}

# Run one SQL text non-interactively. A capped write tees psql's stdout to
# read its verdict line back: ROLLBACK exits 4, UNCAPPED exits 5.
run_sql() {
    local sql="$1"
    refuse_session_set "$sql"
    if ! needs_cap "$sql"; then
        wrap_ro "$sql" | ssh -o IdentityAgent=none "$HOST" "$REMOTE_SCRIPTED"
        return
    fi
    refuse_tx_control "$sql"
    OUT="$(mktemp)"
    trap 'rm -f "$OUT"' EXIT
    wrap_cap "$sql" | ssh -o IdentityAgent=none "$HOST" "$REMOTE_SCRIPTED" | tee "$OUT"
    if grep -q '^prod-psql: ROLLBACK' "$OUT"; then
        echo "prod-psql: rolled back — narrow the WHERE, or set PRECIS_PROD_WRITE_MAX_ROWS for a deliberate bulk change." >&2
        exit 4
    fi
    if grep -q '^prod-psql: UNCAPPED' "$OUT"; then
        exit 5
    fi
}

if [[ $# -gt 0 ]]; then
    # One-shot: the SQL goes over ssh via psql's STDIN, never interpolated into
    # the remote command line — so any quotes/`$`/`;` in the SQL are safe.
    run_sql "$*"
elif [[ ! -t 0 ]]; then
    # Piped SQL on stdin — read it whole so it can be checked first.
    run_sql "$(cat)"
else
    # Interactive: nothing to pre-check; the ⚠ above applies to what you type.
    # With --ro the session is agent_ro, which cannot write at all.
    # No args, interactive terminal → interactive psql (allocate a tty).
    exec ssh -t -o IdentityAgent=none "$HOST" "$REMOTE"
fi
