Repository: zalando/patroni
Pull request #2642 (2023-04-11 14:42:09)

Add before stop hook
We'd like the ability to run scripts prior to shutting down our Postgres instances.

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.

#2596 

CyberDem0n: Thanks for working on it @le-duane.

Here are my thoughts. Currently callbacks are supposed to be fired and forgot and they will be running in parallel with whatever Patroni does with Postgres.
In some cases they are indeed executed only after reaching certain stage:
- on_start is executed when Postgres starts accepting connections
- on_stop is executed when Postgres is already stopped

But, right after `on_stop` there could be `on_start` and therefore the newly executed callback usually cancels the one that already running (Patroni kills the process and its all children).

The `on_pre_stop` feels like a bit special. There is one very useful use-case when we would like to run it not as fire-and-forget, but actually wait for the result. More specifically I am talking about `pgbouncer PAUSE`.

Of course, after promoting the new node (or when Postgres is back running as a primary after a restart) we want to do `pgbouncer RESUME`, but this one we don't need to wait for.

Also, we already have `postgresql.pre_promote` hook, which exit code is actually being checked and `pg_ctl promote` isn't called if the exit code is different from 0.

Having said that I'm thinking that `on_pre_stop` should not be a callback, but something similar to `pre_promote`, because we may want to wait until it finished before executing `pg_ctl stop -mf`. And probably `before_stop` would be a bit better name.
le-duane: Thanks for the quick reply @CyberDem0n - I'll look into `pre_promote`
coveralls: ## Pull Request Test Coverage Report for [Build 4682772676](https://coveralls.io/builds/58850307)

* **5** of **11**   **(45.45%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.05%**) to **99.771%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [patroni/postgresql/__init__.py](https://coveralls.io/builds/58850307/source?filename=patroni%2Fpostgresql%2F__init__.py#L1039) | 5 | 11 | 45.45%
<!-- | **Total:** | **5** | **11** | **45.45%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/58850307/badge)](https://coveralls.io/builds/58850307) |
| :-- | --: |
| Change from base [Build 4619043259](https://coveralls.io/builds/58598229): |  -0.05% |
| Covered Lines: | 12208 |
| Relevant Lines: | 12236 |

---
##### 💛  - [Coveralls](https://coveralls.io)

le-duane: Thanks for the tips @CyberDem0n. I'm a lot happier that this is running synchronously now. 

I changed things a little so I could add a test.
hughcapet: 👍 
CyberDem0n: :+1:

Patch:
diff --git a/docs/SETTINGS.rst b/docs/SETTINGS.rst
index 9390b7c6..bbfef98a 100644
--- a/docs/SETTINGS.rst
+++ b/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:
 
diff --git a/patroni/postgresql/__init__.py b/patroni/postgresql/__init__.py
index c88568a7..1ecf111a 100644
--- a/patroni/postgresql/__init__.py
+++ b/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


Test patch:
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 31a807f5..49e58c7a 100644
--- a/tests/test_postgresql.py
+++ b/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())