Implement kubernetes.bootstrap_labels (patroni#3257)

Allow to define labels that will be assigned to a postgres instance pod when in 'initializing new cluster', 'running custom bootstrap script', 'starting after custom bootstrap', or 'creating replica' state

Author: hughcapet

docs/ENVIRONMENT.rst

@@ -113,6 +113,7 @@ Kubernetes
 -  **PATRONI\_KUBERNETES\_NAMESPACE**: (optional) Kubernetes namespace where the Patroni pod is running. Default value is `default`.
 -  **PATRONI\_KUBERNETES\_LABELS**: Labels in format ``{label1: value1, label2: value2}``. These labels will be used to find existing objects (Pods and either Endpoints or ConfigMaps) associated with the current cluster. Also Patroni will set them on every object (Endpoint or ConfigMap) it creates.
 -  **PATRONI\_KUBERNETES\_SCOPE\_LABEL**: (optional) name of the label containing cluster name. Default value is `cluster-name`.
+-  **PATRONI\_KUBERNETES\_BOOTSTRAP\_LABELS**: (optional) Labels in format ``{label1: value1, label2: value2}``. These labels will be assigned to a Patroni pod when its state is either ``initializing new cluster``, ``running custom bootstrap script``, ``starting after custom bootstrap`` or ``creating replica``.
 -  **PATRONI\_KUBERNETES\_ROLE\_LABEL**: (optional) name of the label containing role (`primary`, `replica` or other custom value). Patroni will set this label on the pod it runs in. Default value is ``role``.
 -  **PATRONI\_KUBERNETES\_LEADER\_LABEL\_VALUE**: (optional) value of the pod label when Postgres role is `primary`. Default value is `primary`.
 -  **PATRONI\_KUBERNETES\_FOLLOWER\_LABEL\_VALUE**: (optional) value of the pod label when Postgres role is `replica`. Default value is `replica`.

docs/yaml_configuration.rst

@@ -178,11 +178,12 @@ Kubernetes
 -  **namespace**: (optional) Kubernetes namespace where Patroni pod is running. Default value is `default`.
 -  **labels**: Labels in format ``{label1: value1, label2: value2}``. These labels will be used to find existing objects (Pods and either Endpoints or ConfigMaps) associated with the current cluster. Also Patroni will set them on every object (Endpoint or ConfigMap) it creates.
 -  **scope\_label**: (optional) name of the label containing cluster name. Default value is `cluster-name`.
+-  **bootstrap\_labels**: (optional) Labels in format ``{label1: value1, label2: value2}``. These labels will be assigned to a Patroni pod when its state is either ``initializing new cluster``, ``running custom bootstrap script``, ``starting after custom bootstrap`` or ``creating replica``.
 -  **role\_label**: (optional) name of the label containing role (`primary`, `replica`, or other custom value). Patroni will set this label on the pod it runs in. Default value is ``role``.
 -  **leader\_label\_value**: (optional) value of the pod label when Postgres role is ``primary``. Default value is ``primary``.
 -  **follower\_label\_value**: (optional) value of the pod label when Postgres role is ``replica``. Default value is ``replica``.
 -  **standby\_leader\_label\_value**: (optional) value of the pod label when Postgres role is ``standby_leader``. Default value is ``primary``.
--  **tmp_\role\_label**: (optional) name of the temporary label containing role (`primary` or `replica`). Value of this label will always use the default of corresponding role. Set only when necessary.
+-  **tmp\_role\_label**: (optional) name of the temporary label containing role (`primary` or `replica`). Value of this label will always use the default of corresponding role. Set only when necessary.
 -  **use\_endpoints**: (optional) if set to true, Patroni will use Endpoints instead of ConfigMaps to run leader elections and keep cluster state.
 -  **pod\_ip**: (optional) IP address of the pod Patroni is running in. This value is required when `use_endpoints` is enabled and is used to populate the leader endpoint subsets when the pod's PostgreSQL is promoted.
 -  **ports**: (optional) if the Service object has the name for the port, the same name must appear in the Endpoint object, otherwise service won't work. For example, if your service is defined as ``{Kind: Service, spec: {ports: [{name: postgresql, port: 5432, targetPort: 5432}]}}``, then you have to set ``kubernetes.ports: [{"name": "postgresql", "port": 5432}]`` and Patroni will use it for updating subsets of the leader Endpoint. This parameter is used only if `kubernetes.use_endpoints` is set.

features/backup_create.py

@@ -3,12 +3,18 @@
 import subprocess
 import sys
 
+from time import sleep
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--datadir", required=True)
     parser.add_argument("--dbname", required=True)
     parser.add_argument("--walmethod", required=True, choices=("fetch", "stream", "none"))
+    parser.add_argument("--sleep", required=False, type=int)
     args, _ = parser.parse_known_args()
 
+    if args.sleep:
+        sleep(args.sleep)
+
     walmethod = ["-X", args.walmethod] if args.walmethod != "none" else []
     sys.exit(subprocess.call(["pg_basebackup", "-D", args.datadir, "-c", "fast", "-d", args.dbname] + walmethod))

features/backup_restore.py

@@ -2,11 +2,17 @@
 import argparse
 import shutil
 
+from time import sleep
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--datadir", required=True)
     parser.add_argument("--sourcedir", required=True)
     parser.add_argument("--test-argument", required=True)
+    parser.add_argument("--sleep", required=False, type=int)
     args, _ = parser.parse_known_args()
 
+    if args.sleep:
+        sleep(args.sleep)
+
     shutil.copytree(args.sourcedir, args.datadir)

features/bootstrap_labels.feature

@@ -0,0 +1,22 @@
+Feature: bootstrap labels
+  Check that user-configurable bootstrap labels are set and removed with state change
+
+Scenario: check label for cluster bootstrap
+    When I start postgres-0
+    Then postgres-0 is a leader after 10 seconds
+    When I start postgres-1 in a cluster batman1 as a long-running clone of postgres-0
+    Then "members/postgres-1" key in DCS has state=running custom bootstrap script after 20 seconds
+    And postgres-1 is labeled with "foo"
+    And postgres-1 is a leader of batman1 after 20 seconds
+
+Scenario: check label for replica bootstrap
+    When I do a backup of postgres-1
+    And I start postgres-2 in cluster batman1 using long-running backup_restore
+    Then "members/postgres-2" key in DCS has state=creating replica after 20 seconds
+    And postgres-2 is labeled with "foo"
+
+Scenario: check bootstrap label is removed
+    Given "members/postgres-1" key in DCS has state=running after 2 seconds
+    And  "members/postgres-2" key in DCS has state=running after 20 seconds
+    Then postgres-1 is not labeled with "foo"
+    And postgres-2 is not labeled with "foo"

features/environment.py

@@ -54,6 +54,8 @@ def start(self, max_wait_limit=5):
         self._log = open(os.path.join(self._output_dir, self._name + '.log'), 'a')
         self._handle = self._start()
 
+        if max_wait_limit < 0:
+            return
         max_wait_limit *= self._context.timeout_multiplier
         for _ in range(max_wait_limit):
             assert self._has_started(), "Process {0} is not running after being started".format(self._name)
@@ -218,6 +220,8 @@ def _make_patroni_test_config(self, name, custom_config):
             'host replication replicator all md5',
             'host all all all md5'
         ]
+        if isinstance(self._context.dcs_ctl, KubernetesController):
+            config['kubernetes'] = {'bootstrap_labels': {'foo': 'bar'}}
 
         if self._context.postgres_supports_ssl and self._context.certfile:
             config['postgresql']['parameters'].update({
@@ -657,6 +661,10 @@ def delete_pod(self, name):
             except Exception:
                 break
 
+    def pod_labels(self, name):
+        pod = self._api.read_namespaced_pod(name, self._namespace)
+        return pod.metadata.labels or {}
+
     def query(self, key, scope='batman', group=None):
         if key.startswith('members/'):
             pod = self._api.read_namespaced_pod(key[8:], self._namespace)
@@ -870,15 +878,16 @@ def create_and_set_output_directory(self, feature_name):
         os.makedirs(feature_dir)
         self._output_dir = feature_dir
 
-    def clone(self, from_name, cluster_name, to_name):
+    def clone(self, from_name, cluster_name, to_name, long_running=False):
         f = self._processes[from_name]
         custom_config = {
             'scope': cluster_name,
             'bootstrap': {
                 'method': 'pg_basebackup',
                 'pg_basebackup': {
                     'command': " ".join(self.BACKUP_SCRIPT
-                                        + ['--walmethod=stream', '--dbname="{0}"'.format(f.backup_source)])
+                                        + ['--walmethod=stream', f'--dbname="{f.backup_source}"',
+                                           f'--sleep {5 if long_running else 0}'])
                 },
                 'dcs': {
                     'postgresql': {
@@ -901,12 +910,16 @@ def clone(self, from_name, cluster_name, to_name):
                 }
             }
         }
-        self.start(to_name, custom_config=custom_config)
+        kwargs = {'custom_config': custom_config}
+        if long_running:
+            kwargs['max_wait_limit'] = -1
+        self.start(to_name, **kwargs)
 
-    def backup_restore_config(self, params=None):
+    def backup_restore_config(self, params=None, long_running=False):
         return {
             'command': (self.BACKUP_RESTORE_SCRIPT
-                        + ' --sourcedir=' + os.path.join(self.patroni_path, 'data', 'basebackup')).replace('\\', '/'),
+                        + ' --sourcedir=' + os.path.join(self.patroni_path, 'data', 'basebackup')
+                        + f' --sleep {5 if long_running else 0}').replace('\\', '/'),
             'test-argument': 'test-value',  # test config mapping approach on custom bootstrap/replica creation
             **(params or {}),
         }
@@ -1126,6 +1139,8 @@ def before_feature(context, feature):
         lib = subprocess.check_output(['pg_config', '--pkglibdir']).decode('utf-8').strip()
         if not os.path.exists(os.path.join(lib, 'citus.so')):
             return feature.skip("Citus extension isn't available")
+    elif feature.name == 'bootstrap labels' and context.dcs_ctl.name() != 'kubernetes':
+        feature.skip("Tested only on Kubernetes")
     context.pctl.create_and_set_output_directory(feature.name)
 
 

features/steps/bootstrap_labels.py

@@ -0,0 +1,31 @@
+from behave import step, then
+
+
+@step('I start {name:name} in a cluster {cluster_name:w} as a long-running clone of {name2:name}')
+def start_cluster_clone(context, name, cluster_name, name2):
+    context.pctl.clone(name2, cluster_name, name, True)
+
+
+@step('I start {name:name} in cluster {cluster_name:w} using long-running backup_restore')
+def start_patroni(context, name, cluster_name):
+    return context.pctl.start(name, custom_config={
+        "scope": cluster_name,
+        "postgresql": {
+            'create_replica_methods': ['backup_restore'],
+            "backup_restore": context.pctl.backup_restore_config(long_running=True),
+            'authentication': {
+                'superuser': {'password': 'patroni1'},
+                'replication': {'password': 'rep-pass1'}
+            }
+        }
+    }, max_wait_limit=-1)
+
+
+@then('{name:name} is labeled with "{label:w}"')
+def pod_labeled(context, name, label):
+    assert label in context.dcs_ctl.pod_labels(name), f'pod {name} is not labeled with {label}'
+
+
+@then('{name:name} is not labeled with "{label:w}"')
+def pod_not_labeled(context, name, label):
+    assert label not in context.dcs_ctl.pod_labels(name), f'pod {name} is still labeled with {label}'

patroni/api.py

@@ -310,12 +310,13 @@ def do_GET(self, write_status_code_only: bool = False) -> None:
         """
         path = '/primary' if self.path == '/' else self.path
         response = self.get_postgresql_status()
+        latest_end_lsn = response.pop('latest_end_lsn', 0)
 
         patroni = self.server.patroni
         cluster = patroni.dcs.cluster
         config = global_config.from_cluster(cluster)
 
-        leader_optime = cluster and cluster.status.last_lsn
+        leader_optime = max(cluster and cluster.status.last_lsn or 0, latest_end_lsn)
         replayed_location = response.get('xlog', {}).get('replayed_location', 0)
         max_replica_lag = parse_int(self.path_query.get('lag', [sys.maxsize])[0], 'B')
         if max_replica_lag is None:
@@ -1298,11 +1299,14 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
 
         :returns: a dict with the status of Postgres/Patroni. The keys are:
 
-            * ``state``: one of :class:`~patroni.postgresql.misc.PostgresqlState` or ``unknown``;
+            * ``state``: Postgres state among ``stopping``, ``stopped``, ``stop failed``, ``crashed``, ``running``,
+              ``starting``, ``start failed``, ``restarting``, ``restart failed``, ``initializing new cluster``,
+              ``initdb failed``, ``running custom bootstrap script``, ``starting after custom bootstrap``,
+              ``custom bootstrap failed``, ``creating replica``, or ``unknown``;
             * ``postmaster_start_time``: ``pg_postmaster_start_time()``;
-            * ``role``: :class:`~patroni.postgresql.misc.PostgresqlRole.REPLICA` or
-                :class:`~patroni.postgresql.misc.PostgresqlRole.PRIMARY` based on ``pg_is_in_recovery()`` output;
+            * ``role``: ``replica`` or ``primary`` based on ``pg_is_in_recovery()`` output;
             * ``server_version``: Postgres version without periods, e.g. ``150002`` for Postgres ``15.2``;
+            * ``latest_end_lsn``: latest_end_lsn value from ``pg_stat_get_wal_receiver()``, only on replica nodes;
             * ``xlog``: dictionary. Its structure depends on ``role``:
 
                 * If :class:`~patroni.postgresql.misc.PostgresqlRole.PRIMARY`:
@@ -1343,15 +1347,18 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
             if postgresql.state not in (PostgresqlState.RUNNING, PostgresqlState.RESTARTING,
                                         PostgresqlState.STARTING):
                 raise RetryFailedError('')
-            replication_state = ('(pg_catalog.pg_stat_get_wal_receiver()).status'
-                                 if postgresql.major_version >= 90600 else 'NULL') + ", " +\
-                ("pg_catalog.current_setting('restore_command')" if postgresql.major_version >= 120000 else "NULL")
+            replication_state = ("pg_catalog.pg_{0}_{1}_diff(wr.latest_end_lsn, '0/0')::bigint, wr.status"
+                                 if postgresql.major_version >= 90600 else "NULL, NULL") + ", " +\
+                ("pg_catalog.current_setting('restore_command')" if postgresql.major_version >= 120000 else "NULL") +\
+                ", " + ("pg_catalog.pg_wal_lsn_diff(wr.written_lsn, '0/0')::bigint"
+                        if postgresql.major_version >= 130000 else "NULL")
             stmt = ("SELECT " + postgresql.POSTMASTER_START_TIME + ", " + postgresql.TL_LSN + ","
                     " pg_catalog.pg_last_xact_replay_timestamp(), " + replication_state + ","
-                    " pg_catalog.array_to_json(pg_catalog.array_agg(pg_catalog.row_to_json(ri))) "
+                    " (SELECT pg_catalog.array_to_json(pg_catalog.array_agg(pg_catalog.row_to_json(ri))) "
                     "FROM (SELECT (SELECT rolname FROM pg_catalog.pg_authid WHERE oid = usesysid) AS usename,"
                     " application_name, client_addr, w.state, sync_state, sync_priority"
-                    " FROM pg_catalog.pg_stat_get_wal_senders() w, pg_catalog.pg_stat_get_activity(pid)) AS ri")
+                    " FROM pg_catalog.pg_stat_get_wal_senders() w, pg_catalog.pg_stat_get_activity(pid)) AS ri)") +\
+                (" FROM pg_catalog.pg_stat_get_wal_receiver() AS wr" if postgresql.major_version >= 90600 else "")
 
             row = self.query(stmt.format(postgresql.wal_name, postgresql.lsn_name,
                                          postgresql.wal_flush), retry=retry)[0]
@@ -1361,7 +1368,7 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
                 'role': PostgresqlRole.REPLICA if row[1] == 0 else PostgresqlRole.PRIMARY,
                 'server_version': postgresql.server_version,
                 'xlog': ({
-                    'received_location': row[4] or row[3],
+                    'received_location': row[10] or row[4] or row[3],
                     'replayed_location': row[3],
                     'replayed_timestamp': row[6],
                     'paused': row[5]} if row[1] == 0 else {
@@ -1383,12 +1390,15 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
                     if not cluster or cluster.is_unlocked() or not cluster.leader else cluster.leader.timeline
                 result['timeline'] = postgresql.replica_cached_timeline(leader_timeline)
 
-            replication_state = postgresql.replication_state_from_parameters(row[1] > 0, row[7], row[8])
+            if row[7]:
+                result['latest_end_lsn'] = row[7]
+
+            replication_state = postgresql.replication_state_from_parameters(row[1] > 0, row[8], row[9])
             if replication_state:
                 result['replication_state'] = replication_state
 
-            if row[9]:
-                result['replication'] = row[9]
+            if row[11]:
+                result['replication'] = row[11]
 
         except (psycopg.Error, RetryFailedError, PostgresConnectionException):
             state = postgresql.state