From 07a9553700310d6d6c2ba5c7e2e4484aeb98d899 Mon Sep 17 00:00:00 2001
From: Konstantin Knizhnik <knizhnik@garret.ru>
Date: Mon, 11 Apr 2022 22:30:08 +0300
Subject: [PATCH] Add test for restore from WAL (#1366)

* Add test for restore from WAL

* Fix python formatting

* Choose unused port in wal restore test

* Move recovery tests to zenith_utils/scripts

* Set LD_LIBRARY_PATH in wal recovery scripts

* Fix python test formatting

* Fix mypy warning

* Bump postgres version

* Bump postgres version
---
 test_runner/batch_others/test_wal_restore.py  | 38 +++++++++++++++++++
 vendor/postgres                               |  2 +-
 zenith_utils/scripts/restore_from_wal.sh      | 20 ++++++++++
 .../scripts/restore_from_wal_archive.sh       | 20 ++++++++++
 4 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 test_runner/batch_others/test_wal_restore.py
 create mode 100755 zenith_utils/scripts/restore_from_wal.sh
 create mode 100755 zenith_utils/scripts/restore_from_wal_archive.sh

diff --git a/test_runner/batch_others/test_wal_restore.py b/test_runner/batch_others/test_wal_restore.py
new file mode 100644
index 0000000000..a5855f2258
--- /dev/null
+++ b/test_runner/batch_others/test_wal_restore.py
@@ -0,0 +1,38 @@
+import os
+import subprocess
+
+from fixtures.utils import mkdir_if_needed
+from fixtures.zenith_fixtures import (ZenithEnvBuilder,
+                                      VanillaPostgres,
+                                      PortDistributor,
+                                      PgBin,
+                                      base_dir,
+                                      vanilla_pg,
+                                      pg_distrib_dir)
+from fixtures.log_helper import log
+
+
+def test_wal_restore(zenith_env_builder: ZenithEnvBuilder,
+                     test_output_dir,
+                     port_distributor: PortDistributor):
+    zenith_env_builder.num_safekeepers = 1
+    env = zenith_env_builder.init_start()
+    env.zenith_cli.create_branch("test_wal_restore")
+    pg = env.postgres.create_start('test_wal_restore')
+    pg.safe_psql("create table t as select generate_series(1,1000000)")
+    tenant_id = pg.safe_psql("show zenith.zenith_tenant")[0][0]
+    env.zenith_cli.pageserver_stop()
+    port = port_distributor.get_port()
+    data_dir = os.path.join(test_output_dir, 'pgsql.restored')
+    restored = VanillaPostgres(data_dir, PgBin(test_output_dir), port)
+    subprocess.call([
+        'bash',
+        os.path.join(base_dir, 'zenith_utils/scripts/restore_from_wal.sh'),
+        os.path.join(pg_distrib_dir, 'bin'),
+        os.path.join(test_output_dir, 'repo/safekeepers/sk1/{}/*'.format(tenant_id)),
+        data_dir,
+        str(port)
+    ])
+    restored.start()
+    assert restored.safe_psql('select count(*) from t') == [(1000000, )]
+    restored.stop()
diff --git a/vendor/postgres b/vendor/postgres
index 8481459996..61afbf978b 160000
--- a/vendor/postgres
+++ b/vendor/postgres
@@ -1 +1 @@
-Subproject commit 848145999653be213141a330569b6f2d9f53dbf2
+Subproject commit 61afbf978b17764134ab6f1650bbdcadac147e71
diff --git a/zenith_utils/scripts/restore_from_wal.sh b/zenith_utils/scripts/restore_from_wal.sh
new file mode 100755
index 0000000000..ef2171312b
--- /dev/null
+++ b/zenith_utils/scripts/restore_from_wal.sh
@@ -0,0 +1,20 @@
+PG_BIN=$1
+WAL_PATH=$2
+DATA_DIR=$3
+PORT=$4
+SYSID=`od -A n -j 24 -N 8 -t d8 $WAL_PATH/000000010000000000000002* | cut -c 3-`
+rm -fr $DATA_DIR
+env -i LD_LIBRARY_PATH=$PG_BIN/../lib $PG_BIN/initdb -E utf8 -D $DATA_DIR --sysid=$SYSID
+echo port=$PORT >> $DATA_DIR/postgresql.conf
+REDO_POS=0x`$PG_BIN/pg_controldata -D $DATA_DIR | fgrep "REDO location"| cut -c 42-`
+declare -i WAL_SIZE=$REDO_POS+114
+$PG_BIN/pg_ctl -D $DATA_DIR -l logfile start
+$PG_BIN/pg_ctl -D $DATA_DIR -l logfile stop -m immediate
+cp $DATA_DIR/pg_wal/000000010000000000000001 .
+cp $WAL_PATH/* $DATA_DIR/pg_wal/
+if [ -f $DATA_DIR/pg_wal/*.partial ]
+then
+	(cd $DATA_DIR/pg_wal ; for partial in \*.partial ; do mv $partial `basename $partial .partial` ; done)
+fi
+dd if=000000010000000000000001 of=$DATA_DIR/pg_wal/000000010000000000000001 bs=$WAL_SIZE count=1 conv=notrunc
+rm -f 000000010000000000000001
diff --git a/zenith_utils/scripts/restore_from_wal_archive.sh b/zenith_utils/scripts/restore_from_wal_archive.sh
new file mode 100755
index 0000000000..07f4fe1e4f
--- /dev/null
+++ b/zenith_utils/scripts/restore_from_wal_archive.sh
@@ -0,0 +1,20 @@
+PG_BIN=$1
+WAL_PATH=$2
+DATA_DIR=$3
+PORT=$4
+SYSID=`od -A n -j 24 -N 8 -t d8 $WAL_PATH/000000010000000000000002* | cut -c 3-`
+rm -fr $DATA_DIR /tmp/pg_wals
+mkdir /tmp/pg_wals
+env -i LD_LIBRARY_PATH=$PG_BIN/../lib $PG_BIN/initdb -E utf8 -U zenith_admin -D $DATA_DIR --sysid=$SYSID
+echo port=$PORT >> $DATA_DIR/postgresql.conf
+REDO_POS=0x`$PG_BIN/pg_controldata -D $DATA_DIR | fgrep "REDO location"| cut -c 42-`
+declare -i WAL_SIZE=$REDO_POS+114
+cp $WAL_PATH/* /tmp/pg_wals
+if [ -f $DATA_DIR/pg_wal/*.partial ]
+then
+	(cd /tmp/pg_wals ; for partial in \*.partial ; do  mv $partial `basename $partial .partial` ; done)
+fi
+dd if=$DATA_DIR/pg_wal/000000010000000000000001 of=/tmp/pg_wals/000000010000000000000001 bs=$WAL_SIZE count=1 conv=notrunc
+echo > $DATA_DIR/recovery.signal
+rm -f $DATA_DIR/pg_wal/*
+echo "restore_command = 'cp /tmp/pg_wals/%f %p'" >> $DATA_DIR/postgresql.conf