Add before stop hook (#2642)

The two cases we have in mind are:
* In spite of following all best practices client-side, logical replication connections can sometimes hang the Postgres shutdown sequence. We'd like to sigterm any misbehaving logical replication connections which remain after x seconds. These will inevitably get killed anyway on master stop timeout.
* remove "role=master" label on current primary when not using k8s as DCS. Waiting until after Postgres fully stops can sometimes be too long for this.
* Pause pgbouncer connections before switchover

Close #2596

docs/SETTINGS.rst

@@ -337,6 +337,7 @@ PostgreSQL
     -  **remove\_data\_directory\_on\_diverged\_timelines**: Patroni will remove the PostgreSQL data directory and recreate the replica if it notices that timelines are diverging and the former primary can not start streaming from the new primary. This option is useful when ``pg_rewind`` can not be used. While performing timelines divergence check on PostgreSQL v10 and older Patroni will try to connect with replication credential to the "postgres" database. Hence, such access should be allowed in the  pg_hba.conf. Default value is **false**.
     -  **replica\_method**: for each create_replica_methods other than basebackup, you would add a configuration section of the same name. At a minimum, this should include "command" with a full path to the actual script to be executed. Other configuration parameters will be passed along to the script in the form "parameter=value".
     -  **pre\_promote**: a fencing script that executes during a failover after acquiring the leader lock but before promoting the replica. If the script exits with a non-zero code, Patroni does not promote the replica and removes the leader key from DCS.
+    -  **before\_stop**: a script that executes immediately prior to stopping postgres. As opposed to a callback, this script runs synchronously, blocking shutdown until it has completed. The return code of this script does not impact whether shutdown proceeds afterwards.
 
 .. _restapi_settings:
 

patroni/postgresql/__init__.py

@@ -729,6 +729,9 @@ class Postgresql(object):
         if not block_callbacks:
             self.set_state('stopping')
 
+        # invoke user-directed before stop script
+        self._before_stop()
+
         if before_shutdown:
             before_shutdown()
 
@@ -1035,6 +1038,21 @@ class Postgresql(object):
             logger.info('pre_promote script `%s` exited with %s', cmd, ret)
         return ret == 0
 
+    def _before_stop(self) -> None:
+        """Synchronously run a script prior to stopping postgres."""
+
+        cmd = self.config.get('before_stop')
+        if cmd:
+            self._do_before_stop(cmd)
+
+    def _do_before_stop(self, cmd: str) -> None:
+        try:
+            ret = self.cancellable.call(shlex.split(cmd))
+            if ret is not None:
+                logger.info('before_stop script `%s` exited with %s', cmd, ret)
+        except Exception as e:
+            logger.error('Exception when calling `%s`: %r', cmd, e)
+
     def promote(self, wait_seconds, task, before_promote=None, on_success=None):
         if self.role in ('promoted', 'master', 'primary'):
             return True

tests/test_postgresql.py

@@ -170,7 +170,8 @@ class TestPostgresql(BaseTestPostgresql):
     @patch('time.sleep', Mock())
     @patch.object(Postgresql, 'is_running')
     @patch.object(Postgresql, '_wait_for_connection_close', Mock())
-    def test_stop(self, mock_is_running):
+    @patch('patroni.postgresql.cancellable.CancellableSubprocess.call')
+    def test_stop(self, mock_cancellable_call, mock_is_running):
         # Postmaster is not running
         mock_callback = Mock()
         mock_is_running.return_value = None
@@ -195,6 +196,19 @@ class TestPostgresql(BaseTestPostgresql):
         mock_postmaster.wait.side_effect = [psutil.TimeoutExpired(30), Mock()]
         self.assertTrue(self.p.stop(on_safepoint=mock_callback, stop_timeout=30))
 
+        # Ensure before_stop script is called when configured to
+        self.p.config._config['before_stop'] = ':'
+        mock_postmaster.wait.side_effect = [psutil.TimeoutExpired(30), Mock()]
+        mock_cancellable_call.return_value = 0
+        with patch('patroni.postgresql.logger.info') as mock_logger:
+            self.p.stop(on_safepoint=mock_callback, stop_timeout=30)
+            self.assertEqual(mock_logger.call_args[0], ('before_stop script `%s` exited with %s', ':', 0))
+        mock_postmaster.wait.side_effect = [psutil.TimeoutExpired(30), Mock()]
+        mock_cancellable_call.side_effect = Exception
+        with patch('patroni.postgresql.logger.error') as mock_logger:
+            self.p.stop(on_safepoint=mock_callback, stop_timeout=30)
+            self.assertEqual(mock_logger.call_args_list[1][0][0], 'Exception when calling `%s`: %r')
+
         # Stop signal failed
         mock_postmaster.signal_stop.return_value = False
         self.assertFalse(self.p.stop())