#!/usr/bin/env bash

export PGUSER=${PGUSER:-postgres}
export PGDATABASE=${PGDATABASE:-postgres}
export SRC_PORT=${SRC_PORT:-5432}
export DST_PORT=${DST_PORT:-5433}

NTABLES=${NTABLES:-100}
NROWS=${NROWS:-10000}

PSQL="$(which psql)"
psql() {
    ${PSQL} -X -v ON_ERROR_STOP=1 "$@"
}

for port in {${SRC_PORT},${DST_PORT}}; do
    for tnum in $(seq 0 ${NTABLES}); do
        echo "CREATE TABLE t${tnum} (id INTEGER PRIMARY KEY);";
    done | psql -p $port;
done

psql -p ${SRC_PORT} -c "CREATE PUBLICATION pub FOR TABLE t0;"
psql -p ${DST_PORT} <<EOF
CREATE SUBSCRIPTION sub
  CONNECTION 'port=${SRC_PORT} user=${PGUSER} dbname=${PGDATABASE}'
  PUBLICATION pub;
EOF

for tnum in $(seq 0 ${NTABLES}); do
    (for id in $(seq 1 ${NROWS}); do echo "INSERT INTO t${tnum} VALUES (${id});"; done | psql -p ${SRC_PORT} -q) &
done

for tnum in $(seq 1 ${NTABLES}); do
    echo "Adding table t${tnum}";
    psql -p ${SRC_PORT} -q -c "ALTER PUBLICATION pub ADD TABLE t${tnum};"
    psql -p ${DST_PORT} -q -c "ALTER SUBSCRIPTION sub REFRESH PUBLICATION;"
    while psql -q 2>/dev/null <<EOF; do
        SELECT 1/0
          FROM pg_subscription_rel AS sr
         INNER JOIN pg_subscription AS s ON (s.oid = srsubid)
         WHERE subname = 'sub'
           AND srrelid = 't${tnum}'::regclass
EOF
        sleep 1;
    done
done

echo "Waiting for data generators ..."
wait

echo "Waiting for replication lag ..."
SRC_LSN=$(psql -p ${SRC_PORT} -q -t -c "SELECT pg_current_wal_lsn()")
while psql -q 2>/dev/null <<EOF; do
    SELECT 1/0
      FROM pg_stat_subscription
     WHERE subname = 'sub'
       AND leader_pid IS NULL AND relid IS NULL
       AND pg_wal_lsn_diff('${SRC_LSN}', latest_end_lsn) <= 0
EOF
    sleep 1;
done

# PG 14 seems to need some extra time, for some reason
sleep 10

echo "Verifying data"
for tnum in $(seq 0 ${NTABLES}); do
    if ! diff -q &>/dev/null \
        <(psql -p ${SRC_PORT} -c "COPY (SELECT * FROM t${tnum} ORDER BY id) TO STDOUT") \
        <(psql -p ${DST_PORT} -c "COPY (SELECT * FROM t${tnum} ORDER BY id) TO STDOUT"); then
        echo "Table t${tnum} has mismatch";
    fi
done