#!/usr/bin/env bash
# Read-only lookups against prod read replicas (IAM auth). IDs/metadata only.
#
#   prod-query.sh <alias> "<select sql>"
#   prod-query.sh <alias> --tables
#   prod-query.sh <alias> --columns <table>
set -euo pipefail

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROFILE="production.rds-reader"
REGION="us-east-1"
DB_USER="reader"
MAX_ROWS=200
DENY_COLS='email|name|phone|birth|dob|address|street|zip|postal|ssn|traits|note|body|content|message|text|diagnos|medical|password|token|secret'

die() { echo "prod-query: $*" >&2; exit 1; }

[[ $# -ge 2 ]] || die "usage: prod-query.sh <alias> \"<select sql>\" | --tables | --columns <table>"

alias_name="$1"; shift
row="$(awk -F'\t' -v a="$alias_name" '$1 == a { print; exit }' "$DIR/databases.tsv")"
[[ -n "$row" ]] || die "unknown alias '$alias_name'. Known: $(cut -f1 "$DIR/databases.tsv" | paste -sd' ' -)"
host="$(cut -f2 <<<"$row")"
database="$(cut -f3 <<<"$row")"

internal=0
case "$1" in
  --tables)
    internal=1
    sql="select table_schema, table_name from information_schema.tables where table_schema not in ('pg_catalog','information_schema') order by 1, 2"
    ;;
  --columns)
    [[ $# -eq 2 ]] || die "--columns needs a table name"
    [[ "$2" =~ ^[A-Za-z0-9_.]+$ ]] || die "invalid table name"
    internal=1
    sql="select column_name, data_type from information_schema.columns where table_name = '${2##*.}' order by ordinal_position"
    ;;
  *)
    sql="$1"
    ;;
esac

if [[ $internal -eq 0 ]]; then
  sql="$(sed -E 's/[[:space:]]*;[[:space:]]*$//' <<<"$sql")"
  lc="$(tr '[:upper:]' '[:lower:]' <<<"$sql" | tr '\n' ' ')"

  [[ "$lc" =~ ^[[:space:]]*(select|with)[[:space:]] ]] || die "only a single SELECT/WITH statement is allowed"
  [[ "$sql" != *";"* ]] || die "multiple statements are not allowed"
  [[ ! "$lc" =~ (^|[^a-z_])(insert|update|delete|drop|alter|create|grant|revoke|truncate|copy|call|do|set|reset|lock|vacuum|pg_read_file|pg_ls_dir|dblink|lo_import|lo_export)([^a-z_]|$) ]] \
    || die "statement contains a disallowed keyword"

  stripped="$(sed -E 's/count\([[:space:]]*\*[[:space:]]*\)//g' <<<"$lc")"
  [[ ! "$stripped" =~ (select|,)[[:space:]]*(distinct[[:space:]]+)?([a-z_0-9]+\.)?\* ]] \
    || die "'*' in the select list is not allowed; list explicit ID/metadata columns (use --columns <table>)"
fi

aws sts get-caller-identity --profile "$PROFILE" >/dev/null 2>&1 \
  || die "no valid AWS SSO session. Run: ! aws sso login --profile $PROFILE"

PGPASSWORD="$(aws --profile "$PROFILE" rds generate-db-auth-token \
  --hostname "$host" --port 5432 --region "$REGION" --username "$DB_USER")"
export PGPASSWORD
export PGHOST="$host" PGPORT=5432 PGDATABASE="$database" PGUSER="$DB_USER" PGSSLMODE=require
export PGOPTIONS="-c default_transaction_read_only=on -c statement_timeout=30000"

out="$(mktemp)"
trap 'rm -f "$out"' EXIT
psql -X --csv -v ON_ERROR_STOP=1 -c "$sql" >"$out"

header="$(head -n 1 "$out")"
if [[ $internal -eq 0 ]]; then
  bad="$(tr ',' '\n' <<<"$header" | tr -d '"' | grep -Ei "$DENY_COLS" || true)"
  if [[ -n "$bad" ]]; then
    echo "prod-query: blocked, output columns look like PII/PHI (rows not shown):" >&2
    echo "$bad" >&2
    echo "Select IDs/metadata only; use these fields only inside WHERE." >&2
    exit 2
  fi
fi

total=$(( $(wc -l <"$out") - 1 ))
head -n $(( MAX_ROWS + 1 )) "$out" \
  | sed -E 's/[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}/<redacted-email>/g'
if (( total > MAX_ROWS )); then
  echo "prod-query: output truncated to $MAX_ROWS of $total rows" >&2
fi