#!/usr/bin/env bash
# Follow the v2/v3 pg_upgrade message literally on an old cluster with
# partitions whose virtual generated column differs from the parent's.
#
# Usage: bash h2.sh.txt OLD=<bindir> FIX=<bindir> TARGET=<bindir> [TARGET2=<bindir>]
#   OLD     unpatched server that creates the divergent partitions
#   FIX     server of the same major that runs the fixes (OLD itself, or
#           the same major with the patch, i.e. after a minor update)
#   TARGET  newer major with the patch: pg_upgrade --check, and the real
#           pg_upgrade at the end; TARGET2 only gets --check
#
# One database per case:
#   case_a  CHECK on the parent that does not involve the column
#   case_b  CHECK on the parent that involves the column; the rows meet
#           it with the partition's expression, not with the parent's
#   case_c  CHECK only on the partition, none on the parent
#   case_d  NOT NULL on the column; a row is NULL with the parent's
#           expression
#   case_e  two levels: only the leaf under a partitioned partition differs
#   case_f  two levels: the partitioned partition differs (and so its leaf)
#
# "Following the message" means only what the text says: for each column
# in the file, take its partitioned table, and run SET EXPRESSION there
# with that table's own expression; on 18, drop the CHECK constraints of
# that table first and add them back afterwards, in one transaction.
set -uo pipefail
for a in "$@"; do eval "${a%%=*}=${a#*=}"; done
: "${OLD:?}" "${FIX:?}" "${TARGET:?}"; TARGET2=${TARGET2:-}
tmp=$(mktemp -d "${TMPDIR:-/tmp}/h2.XXXXXX")
port=${PORT_BASE:-62100}
trap 'for d in "$tmp"/old "$tmp"/new; do [ -f "$d/postmaster.pid" ] && \
  "$FIX/pg_ctl" -D "$d" -m immediate -w stop >/dev/null 2>&1; \
  [ -f "$d/postmaster.pid" ] && "$TARGET/pg_ctl" -D "$d" -m immediate -w stop \
  >/dev/null 2>&1; done; rm -rf "$tmp"' EXIT
major() { "$1/postgres" --version | sed -E 's/.* ([0-9]+).*/\1/'; }
say() { printf '\n### %s\n' "$*"; }
start() { "$1/pg_ctl" -D "$2" -l "$2.log" -w \
            -o "-p $port -k $tmp -c listen_addresses=''" start >/dev/null; }
stop() { "$1/pg_ctl" -D "$2" -m fast -w stop >/dev/null; }
q() { local db=$1; shift; "$FIX/psql" -X -h "$tmp" -p "$port" -U postgres \
        -d "$db" "$@" 2>&1; }
# the SQL is echoed as run, so the output shows each command and its result
run() { local db=$1; shift; "$FIX/psql" -X -e -h "$tmp" -p "$port" \
          -U postgres -d "$db" "$@" 2>&1 | sed 's/^/    /'; }

declare -A SQL
SQL[case_a]="
create table ta (x int check (x > 0),
  v int generated always as (x * 2) virtual) partition by list (x);
create table ta1 partition of ta
  (v generated always as (x * 3) virtual) for values in (1);
insert into ta values (1);"
SQL[case_b]="
create table tb (x int,
  v int generated always as (-x) virtual check (v > 0))
  partition by list (x);
create table tb1 partition of tb
  (v generated always as (x) virtual) for values in (5);
insert into tb values (5);"
SQL[case_c]="
create table tc (x int, v int generated always as (x * 2) virtual)
  partition by list (x);
create table tc1 partition of tc
  (v generated always as (x * 3) virtual) for values in (1);
alter table tc1 add constraint tc1_x_check check (x > 0);
insert into tc values (1);"
SQL[case_d]="
create table td (x int,
  v int generated always as (nullif(x, 0)) virtual not null)
  partition by list (x);
create table td1 partition of td
  (v generated always as (x + 1) virtual) for values in (0);
insert into td values (0);"
SQL[case_e]="
create table te (x int, y int, v int generated always as (x * 2) virtual)
  partition by list (x);
create table te1 partition of te for values in (1) partition by list (y);
create table te1a partition of te1
  (v generated always as (x * 3) virtual) for values in (1);
insert into te values (1, 1);"
SQL[case_f]="
create table tf (x int, y int, v int generated always as (x * 2) virtual)
  partition by list (x);
create table tf1 partition of tf
  (v generated always as (x * 3) virtual) for values in (1)
  partition by list (y);
create table tf1a partition of tf1 for values in (1);
insert into tf values (1, 1);"
CASES="case_a case_b case_c case_d case_e case_f"

check() { # check LABEL BINDIR OLDBIN -> prints pg_upgrade output and the file
  local label=$1 bin=$2 ob=$3 f
  rm -rf "$tmp/new"; "$bin/initdb" -D "$tmp/new" -U postgres -A trust >/dev/null
  echo "  \$ pg_upgrade --check  ($("$ob/postgres" --version) -> $label)"
  ( cd "$tmp" && "$bin/pg_upgrade" --check -U postgres -b "$ob" -B "$bin" \
      -d "$tmp/old" -D "$tmp/new" -p $((port + 1)) -P $((port + 2)) \
      -s "$tmp" ) 2>&1 | grep -vE '^\s*$' | sed -E "s|$tmp|<tmp>|g; s/^/  /"
  f=$(find "$tmp/new" -name partition_virtual_generated_columns.txt | head -1)
  if [ -n "$f" ]; then echo "  contents of partition_virtual_generated_columns.txt:"
    sed 's/^/    /' "$f"; cp "$f" "$tmp/report.txt"
  else rm -f "$tmp/report.txt"; fi
}

say "0. builds"
echo "  OLD    $("$OLD/postgres" --version)"
echo "  FIX    $("$FIX/postgres" --version)"
echo "  TARGET $("$TARGET/postgres" --version)"
[ -n "$TARGET2" ] && echo "  TARGET2 $("$TARGET2/postgres" --version)"

say "1. create the cases with the unpatched server"
"$OLD/initdb" -D "$tmp/old" -U postgres -A trust >/dev/null
start "$OLD" "$tmp/old"
for c in $CASES; do
  "$OLD/createdb" -h "$tmp" -p "$port" -U postgres "$c"
  "$OLD/psql" -X -q -h "$tmp" -p "$port" -U postgres -d "$c" -c "${SQL[$c]}" 2>&1
  echo "  $c: partitions created: $(echo "${SQL[$c]}" | grep -c "partition of")"
done
stop "$OLD" "$tmp/old"

say "2. pg_upgrade --check before any fix"
check TARGET "$TARGET" "$FIX"
[ -n "$TARGET2" ] && check TARGET2 "$TARGET2" "$FIX"
cp "$tmp/report.txt" "$tmp/report-0.txt" 2>/dev/null

say "3. follow the message on $("$FIX/postgres" --version)"
start "$FIX" "$tmp/old"
fixmajor=$(major "$FIX")
db=; add=()
while IFS= read -r line; do
  case $line in
    "In database: "*) db=${line#In database: }; echo; echo "  [$db]"; continue ;;
  esac
  col=${line##*.}; rel=${line%.*}; rel=${rel# }; rel=${rel# }
  parent=$(q "$db" -Atc "select inhparent::regclass from pg_inherits
                          where inhrelid = '$rel'::regclass")
  expr=$(q "$db" -Atc "select pg_get_expr(d.adbin, d.adrelid)
                        from pg_attrdef d join pg_attribute a
                          on a.attrelid = d.adrelid and a.attnum = d.adnum
                       where a.attrelid = '$parent'::regclass
                         and a.attname = '$col'")
  echo "  file lists $rel.$col; its partitioned table is $parent, expression $expr"
  cmds=(-c begin)
  if [ "$fixmajor" = 18 ]; then
    while IFS='|' read -r cn cd; do
      [ -n "$cn" ] || continue
      cmds+=(-c "alter table $parent drop constraint $cn")
      add+=(-c "alter table $parent add constraint $cn $cd")
    done < <(q "$db" -Atc "select conname, pg_get_constraintdef(oid)
                            from pg_constraint
                           where conrelid = '$parent'::regclass
                             and contype = 'c' order by conname")
  fi
  cmds+=(-c "alter table $parent alter column $col set expression as ($expr)")
  cmds+=("${add[@]}"); unset add; add=()
  cmds+=(-c commit)
  run "$db" "${cmds[@]}"
done < "$tmp/report-0.txt"
add=()
stop "$FIX" "$tmp/old"

say "4. pg_upgrade --check after following the message"
check TARGET "$TARGET" "$FIX"
[ -n "$TARGET2" ] && check TARGET2 "$TARGET2" "$FIX"

say "5. what the message does not say: the remaining cases by hand"
start "$FIX" "$tmp/old"
echo "  [case_b] the rows do not meet the CHECK with the parent's expression"
run case_b -c "select x, v as v_via_partition, -x as v_via_parent from tb1"
if [ "$fixmajor" = 18 ]; then
  echo "  [case_c] drop the CHECK of the partition instead, then add it back"
  run case_c -c begin -c "alter table tc1 drop constraint tc1_x_check" \
    -c "alter table tc alter column v set expression as (x * 2)" \
    -c "alter table tc1 add constraint tc1_x_check check (x > 0)" -c commit
fi
echo "  [case_e] SET EXPRESSION on the root instead of the listed partition's parent"
run case_e -c "alter table te alter column v set expression as (x * 2)"
echo "  [case_d] the row is NULL with the parent's expression"
run case_d -c "select x, v as v_via_partition, nullif(x, 0) as v_via_parent from td1"
stop "$FIX" "$tmp/old"
check TARGET "$TARGET" "$FIX"

say "6. real pg_upgrade -> TARGET, without case_b and case_d"
start "$FIX" "$tmp/old"
for c in case_b case_d; do "$FIX/dropdb" -h "$tmp" -p "$port" -U postgres "$c"; done
stop "$FIX" "$tmp/old"
rm -rf "$tmp/new"; "$TARGET/initdb" -D "$tmp/new" -U postgres -A trust >/dev/null
( cd "$tmp" && "$TARGET/pg_upgrade" -U postgres -b "$FIX" -B "$TARGET" \
    -d "$tmp/old" -D "$tmp/new" -p $((port + 1)) -P $((port + 2)) \
    -s "$tmp" ) >"$tmp/upgrade.out" 2>&1
echo "  pg_upgrade exit $?: $(grep -E 'Upgrade Complete|fatal|failure' "$tmp/upgrade.out")"
start "$TARGET" "$tmp/new"
P() { "$TARGET/psql" -X -h "$tmp" -p "$port" -U postgres "$@" 2>&1 | sed 's/^/    /'; }
echo "  each leaf: rows read through the root, rows read directly, its expression"
for t in case_a:ta:ta1 case_c:tc:tc1 case_e:te:te1a case_f:tf:tf1a; do
  IFS=: read -r d r l <<<"$t"
  P -d "$d" -At -c "select '$r/$l', (select string_agg(x||'='||v, ',') from $r),
    (select string_agg(x||'='||v, ',') from $l),
    (select pg_get_expr(adbin, adrelid) from pg_attrdef
      where adrelid = '$l'::regclass and adnum =
        (select attnum from pg_attribute where attrelid = '$l'::regclass
            and attname = 'v'))"
done
echo "  count(*) where v = 2 through each root, constraint_exclusion on / partition"
for t in case_a:ta case_c:tc case_e:te case_f:tf; do
  P -d "${t%%:*}" -At -c "set constraint_exclusion = on" \
    -c "select '${t#*:} on', count(*) from ${t#*:} where v = 2" \
    -c "set constraint_exclusion = partition" \
    -c "select '${t#*:} partition', count(*) from ${t#*:} where v = 2" | grep -v SET
done
echo "  repro-ce-on.sql in a new database of the upgraded cluster"
"$TARGET/createdb" -h "$tmp" -p "$port" -U postgres repro
[ -f "${REPRO:-}" ] && P -d repro -f "$REPRO" | sed -E 's|psql:[^:]*:|psql:repro:|'
stop "$TARGET" "$tmp/new"