Fix recovery slot handling before pg_basebackup

Author: aviralgarg05

gpMgmt/bin/gppylib/commands/pg.py

@@ -1,13 +1,15 @@
 #!/usr/bin/env python3
 #
-# Copyright (c) Greenplum Inc 2008. All Rights Reserved. 
+# Copyright (c) Greenplum Inc 2008. All Rights Reserved.
 #
 
+from contextlib import closing
 import os
 import pipes
 
 from gppylib.gplog import *
 from gppylib.gparray import *
+from gppylib.db import dbconn
 from .base import *
 from .unix import *
 from gppylib.commands.base import *
@@ -17,6 +19,31 @@
 
 GPHOME=os.environ.get('GPHOME')
 
+
+def ensure_replication_slot_exists(source_host, source_port,
+                                   replication_slot_name):
+    if not replication_slot_name:
+        return False
+
+    escaped_slot_name = replication_slot_name.replace("'", "''")
+    dburl = dbconn.DbURL(hostname=source_host, port=source_port,
+                         dbname='template1')
+    with closing(dbconn.connect(dburl, utility=True)) as conn:
+        slot_exists = dbconn.querySingleton(
+            conn,
+            "SELECT count(*) FROM pg_catalog.pg_replication_slots "
+            "WHERE slot_name = '{}'".format(escaped_slot_name))
+        if slot_exists > 0:
+            return False
+
+        dbconn.querySingleton(
+            conn,
+            "SELECT slot_name "
+            "FROM pg_catalog.pg_create_physical_replication_slot('{}')"
+            .format(escaped_slot_name))
+
+    return True
+
 class DbStatus(Command):
     def __init__(self,name,db,ctxt=LOCAL,remoteHost=None):
         self.db=db        

gpMgmt/bin/gppylib/commands/test/unit/test_unit_pg_base_backup.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import Mock, patch
 from gppylib.commands import pg
 
 
@@ -44,6 +45,57 @@ def test_base_backup_does_not_pass_conflicting_xlog_method_argument_when_given_r
         self.assertNotIn("-x", base_backup.command_tokens)
         self.assertNotIn("--xlog", base_backup.command_tokens)
 
+    @patch('gppylib.commands.pg.dbconn.querySingleton', return_value=1)
+    @patch('gppylib.commands.pg.dbconn.connect')
+    @patch('gppylib.commands.pg.dbconn.DbURL')
+    def test_ensure_replication_slot_exists_returns_false_when_slot_exists(self, mock_dburl,
+                                                                           mock_connect,
+                                                                           mock_query_singleton):
+        mock_conn = Mock()
+        mock_connect.return_value = mock_conn
+
+        created = pg.ensure_replication_slot_exists('source-host', 5432, 'slot_name')
+
+        self.assertFalse(created)
+        mock_dburl.assert_called_once_with(hostname='source-host', port=5432, dbname='template1')
+        mock_connect.assert_called_once_with(mock_dburl.return_value, utility=True)
+        self.assertEqual(1, mock_query_singleton.call_count)
+        self.assertIn("FROM pg_catalog.pg_replication_slots", mock_query_singleton.call_args[0][1])
+        mock_conn.close.assert_called_once_with()
+
+    @patch('gppylib.commands.pg.dbconn.querySingleton', side_effect=[0, 'slot_name'])
+    @patch('gppylib.commands.pg.dbconn.connect')
+    @patch('gppylib.commands.pg.dbconn.DbURL')
+    def test_ensure_replication_slot_exists_creates_missing_slot(self, mock_dburl,
+                                                                 mock_connect,
+                                                                 mock_query_singleton):
+        mock_conn = Mock()
+        mock_connect.return_value = mock_conn
+
+        created = pg.ensure_replication_slot_exists('source-host', 5432, 'slot_name')
+
+        self.assertTrue(created)
+        mock_dburl.assert_called_once_with(hostname='source-host', port=5432, dbname='template1')
+        mock_connect.assert_called_once_with(mock_dburl.return_value, utility=True)
+        self.assertEqual(2, mock_query_singleton.call_count)
+        self.assertIn("FROM pg_catalog.pg_replication_slots", mock_query_singleton.call_args_list[0][0][1])
+        self.assertIn("pg_create_physical_replication_slot('slot_name')",
+                      mock_query_singleton.call_args_list[1][0][1])
+        mock_conn.close.assert_called_once_with()
+
+    @patch('gppylib.commands.pg.dbconn.querySingleton')
+    @patch('gppylib.commands.pg.dbconn.connect')
+    @patch('gppylib.commands.pg.dbconn.DbURL')
+    def test_ensure_replication_slot_exists_skips_empty_slot_name(self, mock_dburl,
+                                                                  mock_connect,
+                                                                  mock_query_singleton):
+        created = pg.ensure_replication_slot_exists('source-host', 5432, None)
+
+        self.assertFalse(created)
+        mock_dburl.assert_not_called()
+        mock_connect.assert_not_called()
+        mock_query_singleton.assert_not_called()
+
 
 if __name__ == '__main__':
     unittest.main()

gpMgmt/bin/gppylib/test/unit/test_unit_gpsegrecovery.py

@@ -106,10 +106,12 @@ def setUp(self):
         self.maxDiff = None
         self.mock_logger = Mock(spec=['log', 'info', 'debug', 'error', 'warn', 'exception'])
         self.apply_patches([
+            patch('gpsegrecovery.ensure_replication_slot_exists'),
             patch('gpsegrecovery.start_segment', return_value=Mock()),
             patch('gpsegrecovery.PgBaseBackup.__init__', return_value=None),
             patch('gpsegrecovery.PgBaseBackup.run')
         ])
+        self.mock_ensure_slot = self.get_mock_from_apply_patch('ensure_replication_slot_exists')
         self.mock_pgbasebackup_run = self.get_mock_from_apply_patch('run')
         self.mock_pgbasebackup_init = self.get_mock_from_apply_patch('__init__')
 
@@ -130,6 +132,7 @@ def tearDown(self):
         super(FullRecoveryTestCase, self).tearDown()
 
     def _assert_basebackup_runs(self, expected_init_args):
+        self.mock_ensure_slot.assert_called_once_with('sdw1', 40000, 'internal_wal_replication_slot')
         self.assertEqual(1, self.mock_pgbasebackup_init.call_count)
         self.assertEqual(expected_init_args, self.mock_pgbasebackup_init.call_args)
         self.assertEqual(1, self.mock_pgbasebackup_run.call_count)
@@ -172,66 +175,51 @@ def test_basebackup_run_no_forceoverwrite_passes(self):
         self._assert_basebackup_runs(expected_init_args1)
         self._assert_cmd_passed()
 
-    def test_basebackup_run_one_exception(self):
-        self.mock_pgbasebackup_run.side_effect = [Exception('backup failed once'), Mock()]
+    def test_basebackup_slot_check_exception(self):
+        self.mock_ensure_slot.side_effect = [Exception('slot check failed')]
 
         self.full_recovery_cmd.run()
 
-        expected_init_args1 = call("/data/mirror0", "sdw1", '40000', create_slot=False,
-                                   replication_slot_name='internal_wal_replication_slot',
-                                   forceoverwrite=True, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        expected_init_args2 = call("/data/mirror0", "sdw1", '40000', create_slot=True,
-                                   replication_slot_name='internal_wal_replication_slot',
-                                   forceoverwrite=True, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        self.assertEqual(2, self.mock_pgbasebackup_init.call_count)
-        self.assertEqual([expected_init_args1, expected_init_args2] , self.mock_pgbasebackup_init.call_args_list)
-        self.assertEqual(2, self.mock_pgbasebackup_run.call_count)
-        self.assertEqual([call(validateAfter=True),call(validateAfter=True)], self.mock_pgbasebackup_run.call_args_list)
-        gpsegrecovery.start_segment.assert_called_once_with(self.seg_recovery_info, self.mock_logger, self.era)
-        self._assert_cmd_passed()
+        self.assertEqual(0, self.mock_pgbasebackup_init.call_count)
+        self.assertEqual(0, self.mock_pgbasebackup_run.call_count)
+        self.assertEqual(0, gpsegrecovery.start_segment.call_count)
+        self._assert_cmd_failed('{"error_type": "full", "error_msg": "slot check failed", "dbid": 2, '
+                                '"datadir": "/data/mirror0", "port": 50000, '
+                                '"progress_file": "/tmp/test_progress_file"}')
 
-    def test_basebackup_run_two_exceptions(self):
-        self.mock_pgbasebackup_run.side_effect=[Exception('backup failed once'),
-                                                Exception('backup failed twice')]
+    def test_basebackup_run_exception(self):
+        self.mock_pgbasebackup_run.side_effect=[Exception('backup failed once')]
 
         self.full_recovery_cmd.run()
 
         expected_init_args1 = call("/data/mirror0", "sdw1", '40000', create_slot=False,
                                    replication_slot_name='internal_wal_replication_slot',
                                    forceoverwrite=True, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        expected_init_args2 = call("/data/mirror0", "sdw1", '40000', create_slot=True,
-                                   replication_slot_name='internal_wal_replication_slot',
-                                   forceoverwrite=True, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        self.assertEqual(2, self.mock_pgbasebackup_init.call_count)
-        self.assertEqual([expected_init_args1, expected_init_args2], self.mock_pgbasebackup_init.call_args_list)
-        self.assertEqual(2, self.mock_pgbasebackup_run.call_count)
-        self.assertEqual([call(validateAfter=True),call(validateAfter=True)], self.mock_pgbasebackup_run.call_args_list)
-        self.mock_logger.info.any_call('Running pg_basebackup failed: backup failed once')
-        self.mock_logger.info.assert_called_with("Re-running pg_basebackup, creating the slot this time")
+        self.mock_ensure_slot.assert_called_once_with('sdw1', 40000, 'internal_wal_replication_slot')
+        self.assertEqual(1, self.mock_pgbasebackup_init.call_count)
+        self.assertEqual([expected_init_args1], self.mock_pgbasebackup_init.call_args_list)
+        self.assertEqual(1, self.mock_pgbasebackup_run.call_count)
+        self.assertEqual([call(validateAfter=True)], self.mock_pgbasebackup_run.call_args_list)
         self.assertEqual(0, gpsegrecovery.start_segment.call_count)
-        self._assert_cmd_failed('{"error_type": "full", "error_msg": "backup failed twice", "dbid": 2, ' \
+        self._assert_cmd_failed('{"error_type": "full", "error_msg": "backup failed once", "dbid": 2, ' \
                                 '"datadir": "/data/mirror0", "port": 50000, "progress_file": "/tmp/test_progress_file"}')
 
-    def test_basebackup_run_no_forceoverwrite_two_exceptions(self):
-        self.mock_pgbasebackup_run.side_effect = [Exception('backup failed once'),
-                                                  Exception('backup failed twice')]
+    def test_basebackup_run_no_forceoverwrite_exception(self):
+        self.mock_pgbasebackup_run.side_effect = [Exception('backup failed once')]
         self.full_recovery_cmd.forceoverwrite = False
 
         self.full_recovery_cmd.run()
 
         expected_init_args1 = call("/data/mirror0", "sdw1", '40000', create_slot=False,
                                    replication_slot_name='internal_wal_replication_slot',
                                    forceoverwrite=False, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        # regardless of the passed in value, second call to pg_basebackup will always have forceoverwrite=True
-        expected_init_args2 = call("/data/mirror0", "sdw1", '40000', create_slot=True,
-                                   replication_slot_name='internal_wal_replication_slot',
-                                   forceoverwrite=True, target_gp_dbid=2, progress_file='/tmp/test_progress_file')
-        self.assertEqual(2, self.mock_pgbasebackup_init.call_count)
-        self.assertEqual([expected_init_args1, expected_init_args2], self.mock_pgbasebackup_init.call_args_list)
-        self.assertEqual(2, self.mock_pgbasebackup_run.call_count)
-        self.assertEqual([call(validateAfter=True),call(validateAfter=True)], self.mock_pgbasebackup_run.call_args_list)
+        self.mock_ensure_slot.assert_called_once_with('sdw1', 40000, 'internal_wal_replication_slot')
+        self.assertEqual(1, self.mock_pgbasebackup_init.call_count)
+        self.assertEqual([expected_init_args1], self.mock_pgbasebackup_init.call_args_list)
+        self.assertEqual(1, self.mock_pgbasebackup_run.call_count)
+        self.assertEqual([call(validateAfter=True)], self.mock_pgbasebackup_run.call_args_list)
         self.assertEqual(0, gpsegrecovery.start_segment.call_count)
-        self._assert_cmd_failed('{"error_type": "full", "error_msg": "backup failed twice", "dbid": 2, ' \
+        self._assert_cmd_failed('{"error_type": "full", "error_msg": "backup failed once", "dbid": 2, ' \
                                 '"datadir": "/data/mirror0", "port": 50000, "progress_file": "/tmp/test_progress_file"}')
 
     def test_basebackup_init_exception(self):
@@ -287,7 +275,10 @@ def tearDown(self):
     @patch('gppylib.commands.pg.PgRewind.run')
     @patch('gpsegrecovery.PgBaseBackup.__init__', return_value=None)
     @patch('gpsegrecovery.PgBaseBackup.run')
-    def test_complete_workflow(self, mock_pgbasebackup_run, mock_pgbasebackup_init, mock_pgrewind_run, mock_pgrewind_init):
+    @patch('gpsegrecovery.ensure_replication_slot_exists')
+    def test_complete_workflow(self, mock_ensure_slot, mock_pgbasebackup_run,
+                               mock_pgbasebackup_init, mock_pgrewind_run,
+                               mock_pgrewind_init):
         mix_confinfo = gppylib.recoveryinfo.serialize_list([
             self.full_r1, self.incr_r2])
         sys.argv = ['gpsegrecovery', '-l', '/tmp/logdir', '--era', '{}'.format(self.era), '-c {}'.format(mix_confinfo)]
@@ -301,17 +292,21 @@ def test_complete_workflow(self, mock_pgbasebackup_run, mock_pgbasebackup_init,
         self.assertEqual(1, mock_pgrewind_init.call_count)
         self.assertEqual(1, mock_pgbasebackup_run.call_count)
         self.assertEqual(1, mock_pgbasebackup_init.call_count)
+        mock_ensure_slot.assert_called_once_with('source_hostname1', 6001, 'internal_wal_replication_slot')
         self.assertRegex(gplog.get_logfile(), '/gpsegrecovery.py_\d+\.log')
 
     @patch('gppylib.commands.pg.PgRewind.__init__', return_value=None)
     @patch('gppylib.commands.pg.PgRewind.run')
     @patch('gpsegrecovery.PgBaseBackup.__init__', return_value=None)
     @patch('gpsegrecovery.PgBaseBackup.run')
-    def test_complete_workflow_exception(self, mock_pgbasebackup_run, mock_pgbasebackup_init, mock_pgrewind_run,
+    @patch('gpsegrecovery.ensure_replication_slot_exists')
+    def test_complete_workflow_exception(self, mock_ensure_slot,
+                                         mock_pgbasebackup_run,
+                                         mock_pgbasebackup_init,
+                                         mock_pgrewind_run,
                                          mock_pgrewind_init):
         mock_pgrewind_run.side_effect = [Exception('pg_rewind failed')]
-        mock_pgbasebackup_run.side_effect = [Exception('pg_basebackup failed once'),
-                                             Exception('pg_basebackup failed twice')]
+        mock_pgbasebackup_run.side_effect = [Exception('pg_basebackup failed once')]
         mix_confinfo = gppylib.recoveryinfo.serialize_list([
             self.full_r1, self.incr_r2])
         sys.argv = ['gpsegrecovery', '-l', '/tmp/logdir', '--era={}'.format(self.era), '-c {}'.format(mix_confinfo)]
@@ -322,14 +317,15 @@ def test_complete_workflow_exception(self, mock_pgbasebackup_run, mock_pgbasebac
 
         self.assertCountEqual('[{"error_type": "incremental", "error_msg": "pg_rewind failed", "dbid": 4, "datadir": "target_data_dir4", '
                               '"port": 5004, "progress_file": "/tmp/progress_file4"} , '
-                              '{"error_type": "full", "error_msg": "pg_basebackup failed twice", "dbid": 1,'
+                              '{"error_type": "full", "error_msg": "pg_basebackup failed once", "dbid": 1,'
                               '"datadir": "target_data_dir1", "port": 5001, "progress_file": "/tmp/progress_file1"}]',
                               buf.getvalue().strip())
         self.assertEqual(1, ex.exception.code)
         self.assertEqual(1, mock_pgrewind_run.call_count)
         self.assertEqual(1, mock_pgrewind_init.call_count)
-        self.assertEqual(2, mock_pgbasebackup_run.call_count)
-        self.assertEqual(2, mock_pgbasebackup_init.call_count)
+        self.assertEqual(1, mock_pgbasebackup_run.call_count)
+        self.assertEqual(1, mock_pgbasebackup_init.call_count)
+        mock_ensure_slot.assert_called_once_with('source_hostname1', 6001, 'internal_wal_replication_slot')
         self.assertRegex(gplog.get_logfile(), '/gpsegrecovery.py_\d+\.log')
 
     @patch('recovery_base.gplog.setup_tool_logging')

gpMgmt/bin/lib/gpconfigurenewsegment

@@ -10,7 +10,7 @@ from optparse import Option, OptionGroup, OptionParser, OptionValueError, SUPPRE
 
 from gppylib.gpparseopts import OptParser, OptChecker
 from gppylib.commands.gp import ModifyConfSetting, SegmentStart, SegmentStop
-from gppylib.commands.pg import PgBaseBackup
+from gppylib.commands.pg import PgBaseBackup, ensure_replication_slot_exists
 from gppylib.db import dbconn
 from gppylib.commands import unix
 from gppylib.commands.pg import DbStatus
@@ -134,6 +134,10 @@ class ConfExpSegCmd(Command):
                         self.progressFile = '%s/pg_basebackup.%s.dbid%s.out' % (gplog.get_logger_dir(),
                                                                                 datetime.datetime.today().strftime('%Y%m%d_%H%M%S'),
                                                                                 self.dbid)
+                    ensure_replication_slot_exists(
+                        self.syncWithSegmentHostname,
+                        self.syncWithSegmentPort,
+                        self.replicationSlotName)
                     # Create a mirror based on the primary
                     cmd = PgBaseBackup(target_datadir=self.datadir,
                                        source_host=self.syncWithSegmentHostname,
@@ -149,30 +153,9 @@ class ConfExpSegCmd(Command):
                         self.set_results(CommandResult(0, b'', b'', True, False))
                         if shouldDeleteProgressFile:
                             os.remove(self.progressFile)
-
                     except Exception as e:
-                        #  If the cluster never has mirrors, cmd will fail
-                        #  quickly because the internal slot doesn't exist.
-                        #  Re-run with `create_slot`.
-                        #  GPDB_12_MERGE_FIXME could we check it before? or let
-                        #  pg_basebackup create slot if not exists.
-                        cmd = PgBaseBackup(target_datadir=self.datadir,
-                                           source_host=self.syncWithSegmentHostname,
-                                           source_port=str(self.syncWithSegmentPort),
-                                           create_slot = True,
-                                           replication_slot_name=self.replicationSlotName,
-                                           forceoverwrite=True,
-                                           target_gp_dbid=self.dbid,
-                                           logfile=self.progressFile)
-                        try:
-                            logger.info("Re-running pg_basebackup, creating the slot this time")
-                            cmd.run(validateAfter=True)
-                            self.set_results(CommandResult(0, b'', b'', True, False))
-                            if shouldDeleteProgressFile:
-                                os.remove(self.progressFile)
-                        except Exception as e:
-                            self.set_results(CommandResult(1, b'', str(e).encode(), True, False))
-                            raise
+                        self.set_results(CommandResult(1, b'', str(e).encode(), True, False))
+                        raise
 
                     logger.info("Successfully ran pg_basebackup: %s" % cmd.cmdStr)
                     return