Merge pull request #460 from aiven/rdunklau-add_walreceiver_tests

Add walreceiver tests

#460

Author: kmichel-aiven

test/conftest.py

@@ -43,6 +43,7 @@ def __init__(self, pgdata):
         self.pgdata = pgdata
         self.pg = None
         self.user = None
+        self._connection = None
 
     @property
     def pgver(self):
@@ -84,6 +85,11 @@ def kill(self, force=True, immediate=True):
         if not force and self.pg.poll() is None:
             raise Exception("PG pid {} not dead".format(self.pg.pid))
 
+    def connect(self):
+        connection_string = pgutil.create_connection_string(self.user)
+        conn = psycopg2.connect(connection_string)
+        return conn
+
 
 @contextlib.contextmanager
 def setup_pg():
@@ -200,6 +206,13 @@ def pghoard(db, tmpdir, request):  # pylint: disable=redefined-outer-name
     yield from pghoard_base(db, tmpdir, request)
 
 
+@pytest.yield_fixture
+def pghoard_walreceiver(db, tmpdir, request):
+    # Initialize with only one transfer agent, as we want a reliable
+    # last transfered state.
+    yield from pghoard_base(db, tmpdir, request, active_backup_mode="walreceiver", transfer_count=1, compression_count=1)
+
+
 @pytest.yield_fixture  # pylint: disable=redefined-outer-name
 def pghoard_separate_volume(db, tmpdir, request):
     tmpfs_volume = os.path.join(str(tmpdir), "tmpfs")
@@ -241,20 +254,20 @@ def pghoard_base(
     metrics_cfg=None,
     *,
     backup_location=None,
-    pg_receivexlog_config=None
+    pg_receivexlog_config=None,
+    active_backup_mode="pg_receivexlog",
+    slot_name=None,
+    compression_count=None
 ):
     test_site = request.function.__name__
 
-    if pg_receivexlog_config:
-        active_backup_mode = "pg_receivexlog"
-    elif os.environ.get("pghoard_test_walreceiver"):
-        active_backup_mode = "walreceiver"
-    else:
-        active_backup_mode = "pg_receivexlog"
-
     if compression == "snappy" and not snappy:
         compression = "lzma"
 
+    node = db.user.copy()
+    if slot_name is not None:
+        node["slot"] = slot_name
+
     backup_location = backup_location or os.path.join(str(tmpdir), "backupspool")
     config = {
         "alert_file_dir": os.path.join(str(tmpdir), "alerts"),
@@ -267,7 +280,7 @@ def pghoard_base(
                 "pg_bin_directory": db.pgbin,
                 "pg_data_directory": db.pgdata,
                 "pg_receivexlog": pg_receivexlog_config or {},
-                "nodes": [db.user],
+                "nodes": [node],
                 "object_storage": {
                     "storage_type": "local",
                     "directory": os.path.join(str(tmpdir), "backups"),
@@ -297,6 +310,9 @@ def pghoard_base(
     if transfer_count is not None:
         config["transfer"] = {"thread_count": transfer_count}
 
+    if compression_count is not None:
+        config["compression"]["thread_count"] = compression_count
+
     confpath = os.path.join(str(tmpdir), "config.json")
     with open(confpath, "w") as fp:
         json.dump(config, fp)

test/test_basebackup.py

@@ -16,14 +16,15 @@
 import psycopg2
 import pytest
 
-from pghoard import common, metrics, pgutil
+from pghoard import common, metrics
 from pghoard.basebackup import PGBaseBackup
 from pghoard.common import BaseBackupMode
 from pghoard.restore import Restore, RestoreError
 from pghoard.rohmu import get_transfer
 from pghoard.rohmu.compat import makedirs
 
 from .conftest import PGTester
+from .util import switch_wal
 
 Restore.log_tracebacks = True
 
@@ -378,10 +379,9 @@ def test_basebackups_local_tar_legacy(self, capsys, db, pghoard, tmpdir):
     def test_basebackups_local_tar_exclusive_conflict(self, capsys, db, pghoard, tmpdir):
         if db.pgver >= "9.6":
             pytest.skip("PostgreSQL < 9.6 required for exclusive backup tests")
-        conn_str = pgutil.create_connection_string(db.user)
         need_stop = False
         try:
-            with psycopg2.connect(conn_str) as conn:
+            with db.connect() as conn:
                 conn.autocommit = True
                 cursor = conn.cursor()
                 cursor.execute("SELECT pg_start_backup('conflicting')")  # pylint: disable=used-before-assignment
@@ -390,14 +390,13 @@ def test_basebackups_local_tar_exclusive_conflict(self, capsys, db, pghoard, tmp
             need_stop = False
         finally:
             if need_stop:
-                with psycopg2.connect(conn_str) as conn:
+                with db.connect() as conn:
                     conn.autocommit = True
                     cursor = conn.cursor()
                     cursor.execute("SELECT pg_stop_backup()")
 
     def test_basebackups_local_tar_pgespresso(self, capsys, db, pghoard, tmpdir):
-        conn_str = pgutil.create_connection_string(db.user)
-        with psycopg2.connect(conn_str) as conn:
+        with db.connect() as conn:
             conn.autocommit = True
             cursor = conn.cursor()
             cursor.execute("SELECT 1 FROM pg_available_extensions WHERE name = 'pgespresso' AND default_version >= '1.2'")
@@ -415,8 +414,7 @@ def test_basebackups_replica_local_tar_nonexclusive(self, capsys, recovery_db, p
         self._test_basebackups(capsys, recovery_db, pghoard, tmpdir, BaseBackupMode.local_tar, replica=True)
 
     def test_basebackups_replica_local_tar_pgespresso(self, capsys, recovery_db, pghoard, tmpdir):
-        conn_str = pgutil.create_connection_string(recovery_db.user)
-        with psycopg2.connect(conn_str) as conn:
+        with recovery_db.connect() as conn:
             conn.autocommit = True
             cursor = conn.cursor()
             cursor.execute("SELECT 1 FROM pg_available_extensions WHERE name = 'pgespresso' AND default_version >= '1.2'")
@@ -433,8 +431,7 @@ def test_basebackups_tablespaces(self, capsys, db, pghoard, tmpdir):
         # tablespace could interfere with other tests
         tspath = tmpdir.join("extra-ts").strpath
         os.makedirs(tspath)
-        conn_str = pgutil.create_connection_string(db.user)
-        conn = psycopg2.connect(conn_str)
+        conn = db.connect()
         conn.autocommit = True
         cursor = conn.cursor()
         cursor.execute("CREATE TABLESPACE tstest LOCATION %s", [tspath])
@@ -451,19 +448,9 @@ def test_basebackups_tablespaces(self, capsys, db, pghoard, tmpdir):
             wal_directory = os.path.join(pghoard.config["backup_location"], pghoard.test_site, "xlog_incoming")
             makedirs(wal_directory, exist_ok=True)
             pghoard.receivexlog_listener(pghoard.test_site, db.user, wal_directory)
-            if conn.server_version >= 100000:
-                cursor.execute("SELECT txid_current(), pg_switch_wal()")
-            else:
-                cursor.execute("SELECT txid_current(), pg_switch_xlog()")
-
+            switch_wal(conn)
             self._test_create_basebackup(capsys, db, pghoard, "local-tar")
-
-            if conn.server_version >= 100000:
-                cursor.execute("SELECT txid_current(), pg_switch_wal()")
-                cursor.execute("SELECT txid_current(), pg_switch_wal()")
-            else:
-                cursor.execute("SELECT txid_current(), pg_switch_xlog()")
-                cursor.execute("SELECT txid_current(), pg_switch_xlog()")
+            switch_wal(conn)
 
             backup_out = tmpdir.join("test-restore").strpath
             backup_ts_out = tmpdir.join("test-restore-tstest").strpath
@@ -574,13 +561,12 @@ def test_basebackups_tablespaces(self, capsys, db, pghoard, tmpdir):
             r_db = PGTester(backup_out)
             r_db.user = dict(db.user, host=backup_out)
             r_db.run_pg()
-            r_conn_str = pgutil.create_connection_string(r_db.user)
 
             # Wait for PG to start up
             start_time = time.monotonic()
             while True:
                 try:
-                    r_conn = psycopg2.connect(r_conn_str)
+                    r_conn = r_db.connect()
                     break
                 except psycopg2.OperationalError as ex:
                     if "starting up" in str(ex):

test/test_pghoard.py

@@ -13,8 +13,6 @@
 from pathlib import Path
 from unittest.mock import Mock, patch
 
-import psycopg2
-
 from pghoard import common
 from pghoard.common import (BaseBackupFormat, create_alert_file, delete_alert_file, write_json_file)
 from pghoard.pghoard import PGHoard
@@ -23,6 +21,7 @@
 from pghoard.rohmu import rohmufile
 
 from .base import PGHoardTestCase
+from .util import switch_wal, wait_for_xlog
 
 
 class TestPGHoard(PGHoardTestCase):
@@ -551,10 +550,8 @@ def clean_alert_files():
 
     def test_pause_on_disk_full(self, db, pghoard_separate_volume, caplog):
         pghoard = pghoard_separate_volume
-        conn_str = create_connection_string(db.user)
-        conn = psycopg2.connect(conn_str)
+        conn = db.connect()
         conn.autocommit = True
-        cursor = conn.cursor()
 
         wal_directory = os.path.join(pghoard.config["backup_location"], pghoard.test_site, "xlog_incoming")
         os.makedirs(wal_directory, exist_ok=True)
@@ -566,70 +563,26 @@ def test_pause_on_disk_full(self, db, pghoard_separate_volume, caplog):
         for _ in range(16):
             # Note: do not combine two function call in one select, PG executes it differently and
             # sometimes looks like it generates less WAL files than we wanted
-            cursor.execute("SELECT txid_current()")
-            if conn.server_version >= 100000:
-                cursor.execute("SELECT pg_switch_wal()")
-            else:
-                cursor.execute("SELECT pg_switch_xlog()")
-
-        start = time.monotonic()
-        while True:
-            xlogs = None
-            # At the start, this is not yet defined
-            transfer_agent_state_for_site = pghoard.transfer_agent_state.get(pghoard.test_site)
-            if transfer_agent_state_for_site:
-                xlogs = transfer_agent_state_for_site["upload"]["xlog"]["xlogs_since_basebackup"]
-                if xlogs >= 15:
-                    break
-
-            if time.monotonic() - start > 15:
-                assert False, "Expected at least 15 xlog uploads, got {}".format(xlogs)
-
-            time.sleep(0.1)
+            switch_wal(conn)
+        conn.close()
+
+        wait_for_xlog(pghoard, 15)
         assert "pausing pg_receive(wal|xlog)" in caplog.text
 
     def test_surviving_pg_receivewal_hickup(self, db, pghoard):
-        conn_str = create_connection_string(db.user)
-        conn = psycopg2.connect(conn_str)
-        conn.autocommit = True
-        cursor = conn.cursor()
-
         wal_directory = os.path.join(pghoard.config["backup_location"], pghoard.test_site, "xlog_incoming")
         os.makedirs(wal_directory, exist_ok=True)
 
-        def trigger_new_wal():
-            # Note: do not combine two function call in one select, PG executes it differently and
-            # sometimes looks like it generates less WAL files than we wanted
-            cursor.execute("SELECT txid_current()")
-            if conn.server_version >= 100000:
-                cursor.execute("SELECT pg_switch_wal()")
-            else:
-                cursor.execute("SELECT pg_switch_xlog()")
-
-        def wait_for_xlog(count: int):
-            start = time.monotonic()
-            while True:
-                xlogs = None
-                # At the start, this is not yet defined
-                transfer_agent_state_for_site = pghoard.transfer_agent_state.get(pghoard.test_site)
-                if transfer_agent_state_for_site:
-                    xlogs = transfer_agent_state_for_site["upload"]["xlog"]["xlogs_since_basebackup"]
-                    if xlogs >= count:
-                        break
-
-                if time.monotonic() - start > 15:
-                    assert False, "Expected at least {} xlog uploads, got {}".format(count, xlogs)
-
-                time.sleep(0.1)
-
         pghoard.receivexlog_listener(pghoard.test_site, db.user, wal_directory)
+        conn = db.connect()
+        conn.autocommit = True
 
         # Make sure we have already a few files so pg_receivewal has something to start from when it eventually restarts
         # +1: to finish the current one
         for _ in range(3 + 1):
-            trigger_new_wal()
+            switch_wal(conn)
 
-        wait_for_xlog(3)
+        wait_for_xlog(pghoard, 3)
 
         # stop pg_receivewal so we cannot process new WAL segments
         pghoard.receivexlogs[pghoard.test_site].running = False
@@ -641,11 +594,12 @@ def wait_for_xlog(count: int):
 
         # add more WAL segments
         for _ in range(10):
-            trigger_new_wal()
+            switch_wal(conn)
+        conn.close()
 
         # restart
         pghoard.receivexlog_listener(pghoard.test_site, db.user, wal_directory)
         assert pghoard.receivexlogs[pghoard.test_site].is_alive()
 
         # We should now process all created segments, not only the ones which were created after pg_receivewal was restarted
-        wait_for_xlog(n_xlogs + 10)
+        wait_for_xlog(pghoard, n_xlogs + 10)

test/test_walreceiver.py

@@ -0,0 +1,66 @@
+from pghoard.pghoard import PGHoard
+from pghoard.wal import (get_current_wal_file, get_previous_wal_on_same_timeline, name_for_tli_log_seg, name_to_tli_log_seg)
+
+from .util import switch_wal, wait_for_xlog
+
+
+def get_transfer_agent_upload_xlog_state(pghoard: PGHoard):
+    transfer_agent_state = pghoard.transfer_agent_state.get(pghoard.test_site)
+    if transfer_agent_state is None:
+        return {}
+    return transfer_agent_state["upload"]["xlog"]
+
+
+def stop_walreceiver(pghoard: PGHoard):
+    walreceiver = pghoard.walreceivers.pop(pghoard.test_site)
+    walreceiver.running = False
+    walreceiver.join()
+    return walreceiver.last_flushed_lsn
+
+
+class TestWalReceiver:
+    def test_walreceiver(self, db, pghoard_walreceiver):
+        """
+        Test the happy-path of the wal receiver.
+        """
+        conn = db.connect()
+        conn.autocommit = True
+        server_version = conn.server_version
+
+        pghoard = pghoard_walreceiver
+        node = pghoard.config["backup_sites"][pghoard.test_site]["nodes"][0]
+        # The transfer agent state will be used to check what
+        # was uploaded
+        # Before starting the walreceiver, get the current wal name.
+        wal_name = get_current_wal_file(node)
+        # Start streaming, force a wal rotation, and check the wal has been
+        # archived
+        pghoard.start_walreceiver(pghoard.test_site, node, None)
+        switch_wal(conn)
+        # Check that we uploaded one file, and it is the right one.
+        wait_for_xlog(pghoard, 1)
+        last_flushed_lsn = stop_walreceiver(pghoard)
+        # Record the last flushed lsn
+        state = get_transfer_agent_upload_xlog_state(pghoard)
+        assert state.get("xlogs_since_basebackup") == 1
+        assert state.get("latest_filename") == wal_name
+
+        # Generate some more wal while the walreceiver is not running,
+        # and check that we can fetch it once done using the recorded state
+        for _ in range(3):
+            switch_wal(conn)
+        conn.close()
+        # The last wal file is the previous one, as the current one is not
+        # complete.
+        wal_name = get_current_wal_file(node)
+        tli, log, seg = name_to_tli_log_seg(wal_name)
+        seg, log = get_previous_wal_on_same_timeline(seg, log, server_version)
+
+        previous_wal_name = name_for_tli_log_seg(tli, log, seg)
+
+        pghoard.start_walreceiver(pghoard.test_site, node, last_flushed_lsn)
+        wait_for_xlog(pghoard, 4)
+        last_flushed_lsn = stop_walreceiver(pghoard)
+        state = get_transfer_agent_upload_xlog_state(pghoard)
+        assert state.get("xlogs_since_basebackup") == 4
+        assert state.get("latest_filename") == previous_wal_name