#!/usr/bin/env bash
set -euo pipefail
. /var/lib/postgresql/scripts/pg/pg.sh

# This is a 'golden' test. It verifies that we can restore to the latest
# committed transaction after an abrupt system crash.
# It also checks that the WAL files generated are byte-for-byte identical to
# those generated by 'pg_receivewal'.

# Steps:
#
# * Initialize and start a PostgreSQL cluster
# * Run WAL receivers (our WAL archiver and pg_receivewal)
# * Create a base backup
# * Create a table, and insert the current timestamp every second (in the background)
# * Run pgbench to populate the database with 1 million rows
# * Generate additional data (~512 MiB)
# * Concurrently create 100 tables with 10000 rows each.
# * Terminate the insert-script job
# * Run pg_dumpall and save the output as plain SQL
# * Terminate all PostgreSQL processes and delete the PGDATA directory
# * Restore PGDATA from the base backup, add recovery.signal, and configure restore_command
# * Rename all *.partial WAL files in the WAL archive directories
# * Start the PostgreSQL cluster (cluster should recover to the latest committed transaction)
# * Run pg_dumpall after the cluster is ready
# * Diff the pg_dumpall results (before and after)
# * Check the insert-script logs and verify that the table contains the last inserted row
# * Compare WAL directories (filenames and contents must match 100%)
# * Clean up WAL directories and rerun the WAL archivers on a new timeline (cleanup is necessary since we run receivers with --no-loop option)
# * Compare the WAL directories again

export BASEBACKUP_PATH="/tmp/basebackup"
export WAL_PATH="/tmp/wal-archive"
export WAL_PATH_PG_RECEIVEWAL="/tmp/wal-archive-pg_receivewal"
export LOG_FILE="/tmp/pgrwl.log"

# Default environment
export PGHOST="localhost"
export PGPORT="5432"
export PGUSER="postgres"
export PGPASSWORD="postgres"

x_remake_dirs() {
  # cleanup possible state
  rm -rf "${BASEBACKUP_PATH}" && mkdir -p "${BASEBACKUP_PATH}"
  rm -rf "${WAL_PATH}" && mkdir -p "${WAL_PATH}"
  rm -rf "${WAL_PATH_PG_RECEIVEWAL}" && mkdir -p "${WAL_PATH_PG_RECEIVEWAL}"

  cat <<EOF >/tmp/config.json
{
  "main": {
    "listen_port": 7070,
    "directory": "/tmp/wal-archive"
  },
  "receiver": {
    "slot": "pgrwl_v5",
    "no_loop": true
  },
  "log": {
    "level": "trace",
    "format": "text",
    "add_source": true
  }
}
EOF
}

x_backup_restore() {
  x_remake_dirs

  # rerun the cluster
  echo_delim "init and run a cluster"
  xpg_rebuild
  xpg_start

  # start receiving from the fresh WAL that is identical for both receivers
  psql -U postgres -c 'drop table if exists xxx; select pg_switch_wal(); create table if not exists xxx (id serial);'

  # run wal-receivers
  echo_delim "running wal-receivers"
  # run wal-receiver
  nohup /usr/local/bin/pgrwl start -c "/tmp/config.json" -m receive >>"$LOG_FILE" 2>&1 &
  # run pg_receivewal
  bash "/var/lib/postgresql/scripts/pg/run_pg_receivewal.sh" "start"

  # make a basebackup before doing anything
  echo_delim "creating basebackup"
  pg_basebackup \
    --pgdata="${BASEBACKUP_PATH}/data" \
    --wal-method=none \
    --checkpoint=fast \
    --progress \
    --no-password \
    --verbose

  # run inserts in a background
  bash "/var/lib/postgresql/scripts/pg/run_inserts.sh" start

  # fill with 1M rows
  echo_delim "running pgbench"
  pgbench -i -s 10 postgres

  # wait a little
  sleep 5

  echo_delim "generate 512Mi of data"
  psql -v ON_ERROR_STOP=1 -U postgres <<-EOSQL
  -- 1Gi/2
  CREATE TABLE bigdata AS
  SELECT i, repeat('x', 1024) AS filler
  FROM generate_series(1, (1 * 1024 * 1024)/2) AS t(i);
EOSQL

  # concurrently populate with random tables+data
  echo_delim "concurrently populate with 100 tables, 10K rows each"
  bash "/var/lib/postgresql/scripts/gendata/gendata.sh" >/dev/null 2>&1

  # stop inserts
  pkill -f inserts.sh

  # remember the state
  pg_dumpall -f /tmp/pg_dumpall-before

  # stop cluster, cleanup data
  echo_delim "teardown"
  xpg_teardown

  # restore from backup
  echo_delim "restoring backup"
  mv "${BASEBACKUP_PATH}/data" "${PGDATA}"
  chmod 0750 "${PGDATA}"
  chown -R postgres:postgres "${PGDATA}"
  touch "${PGDATA}/recovery.signal"

  # prepare archive (all partial files contain valid wal-segments)
  find "${WAL_PATH}" -type f -name "*.partial" -exec bash -c 'for f; do mv -v "$f" "${f%.partial}"; done' _ {} +
  find "${WAL_PATH_PG_RECEIVEWAL}" -type f -name "*.partial" -exec bash -c 'for f; do mv -v "$f" "${f%.partial}"; done' _ {} +

  # fix configs
  xpg_config
  cat <<EOF >>"${PG_CFG}"
#restore_command = 'cp ${WAL_PATH}/%f %p'
restore_command = 'pgrwl restore-command --serve-addr=127.0.0.1:7070 %f %p'
EOF

  # run serve-mode
  echo_delim "running wal fetcher"
  nohup /usr/local/bin/pgrwl start -c "/tmp/config.json" -m serve >>"$LOG_FILE" 2>&1 &

  # cleanup logs
  >/var/log/postgresql/pg.log

  # run restored cluster
  echo_delim "running cluster"
  xpg_start

  # wait until is in recovery, check logs, etc...
  xpg_wait_is_in_recovery
  cat /var/log/postgresql/pg.log

  # check diffs
  echo_delim "running diff on pg_dumpall dumps (before vs after)"
  pg_dumpall -f /tmp/pg_dumpall-arter
  diff /tmp/pg_dumpall-before /tmp/pg_dumpall-arter

  # read the latest rec
  echo_delim "read latest applied records"
  echo "table content:"
  psql --pset pager=off -c "select * from public.tslog;"
  echo "insert log content:"
  tail -10 /tmp/insert-ts.log

  # compare with pg_receivewal
  echo_delim "compare wal-archive with pg_receivewal"
  find "${WAL_PATH}" -type f -name "*.json" -delete
  bash "/var/lib/postgresql/scripts/utils/dircmp.sh" "${WAL_PATH}" "${WAL_PATH_PG_RECEIVEWAL}"

  # run receivers with a new timeline
  echo_delim "cleanup wal-archives, run wal-receivers with a new timeline"
  x_remake_dirs
  # run wal-receiver
  pkill -f pgrwl || true
  nohup /usr/local/bin/pgrwl start -c "/tmp/config.json" -m receive >>"$LOG_FILE" 2>&1 &
  # run pg_receivewal
  bash "/var/lib/postgresql/scripts/pg/run_pg_receivewal.sh" "start"
  sleep 5 # wait a little (feedback timeout, etc...)

  # compare with pg_receivewal
  echo_delim "compare wal-archive with pg_receivewal with a new timeline stream"
  find "${WAL_PATH}" -type f -name "*.json" -delete
  bash "/var/lib/postgresql/scripts/utils/dircmp.sh" "${WAL_PATH}" "${WAL_PATH_PG_RECEIVEWAL}"
}

x_backup_restore "${@}"
