Implement kubernetes.bootstrap_labels (#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

docs/ENVIRONMENT.rst

@@ -117,6 +117,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

@@ -182,11 +182,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 argparse
 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 @@ class AbstractController(abc.ABC):
         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 @@ class PatroniController(AbstractController):
             '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 @@ class KubernetesController(AbstractExternalDcsController):
             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,7 +878,7 @@ class PatroniPoolController(object):
         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,
@@ -878,7 +886,8 @@ class PatroniPoolController(object):
                 '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 @@ class PatroniPoolController(object):
                 }
             }
         }
-        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 {}),
         }
@@ -1124,6 +1137,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

@@ -1265,8 +1265,8 @@ class RestApiHandler(BaseHTTPRequestHandler):
 
             * ``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``, ``custom bootstrap failed``,
-              ``creating replica``, or ``unknown``;
+              ``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``: ``replica`` or ``primary`` based on ``pg_is_in_recovery()`` output;
             * ``server_version``: Postgres version without periods, e.g. ``150002`` for Postgres ``15.2``;

patroni/config.py

@@ -663,7 +663,7 @@ class Config(object):
                               'SERVICE_TAGS', 'NAMESPACE', 'CONTEXT', 'USE_ENDPOINTS', 'SCOPE_LABEL', 'ROLE_LABEL',
                               'POD_IP', 'PORTS', 'LABELS', 'BYPASS_API_SERVICE', 'RETRIABLE_HTTP_CODES', 'KEY_PASSWORD',
                               'USE_SSL', 'SET_ACLS', 'GROUP', 'DATABASE', 'LEADER_LABEL_VALUE', 'FOLLOWER_LABEL_VALUE',
-                              'STANDBY_LEADER_LABEL_VALUE', 'TMP_ROLE_LABEL', 'AUTH_DATA') and name:
+                              'STANDBY_LEADER_LABEL_VALUE', 'TMP_ROLE_LABEL', 'AUTH_DATA', 'BOOTSTRAP_LABELS') and name:
                     value = os.environ.pop(param)
                     if name == 'CITUS':
                         if suffix == 'GROUP':
@@ -674,7 +674,7 @@ class Config(object):
                         value = value and parse_int(value)
                     elif suffix in ('HOSTS', 'PORTS', 'CHECKS', 'SERVICE_TAGS', 'RETRIABLE_HTTP_CODES'):
                         value = value and _parse_list(value)
-                    elif suffix in ('LABELS', 'SET_ACLS', 'AUTH_DATA'):
+                    elif suffix in ('LABELS', 'SET_ACLS', 'AUTH_DATA', 'BOOTSTRAP_LABELS'):
                         value = _parse_dict(value)
                     elif suffix in ('USE_PROXIES', 'REGISTER_SERVICE', 'USE_ENDPOINTS', 'BYPASS_API_SERVICE', 'VERIFY'):
                         value = parse_bool(value)

patroni/ctl.py

@@ -1535,8 +1535,8 @@ def output_members(cluster: Cluster, name: str, extended: bool = False,
         * ``Role``: ``Leader``, ``Standby Leader``, ``Sync Standby`` or ``Replica``;
         * ``State``: ``stopping``, ``stopped``, ``stop failed``, ``crashed``, ``running``, ``starting``,
           ``start failed``, ``restarting``, ``restart failed``, ``initializing new cluster``, ``initdb failed``,
-          ``running custom bootstrap script``, ``custom bootstrap failed``, ``creating replica``, ``streaming``,
-          ``in archive recovery``, and so on;
+          ``running custom bootstrap script``, ``starting after custom bootstrap``, ``custom bootstrap failed``,
+          ``creating replica``, ``streaming``, ``in archive recovery``, and so on;
         * ``TL``: current timeline in Postgres;
           ``Lag in MB``: replication lag.
 

patroni/dcs/kubernetes.py

@@ -760,6 +760,8 @@ class Kubernetes(AbstractDCS):
         self._follower_label_value = config.get('follower_label_value', 'replica')
         self._standby_leader_label_value = config.get('standby_leader_label_value', 'primary')
         self._tmp_role_label = config.get('tmp_role_label')
+        self._bootstrap_labels: Dict[str, str] = {str(k): str(v)
+                                                  for k, v in (config.get('bootstrap_labels') or EMPTY_DICT).items()}
         self._ca_certs = os.environ.get('PATRONI_KUBERNETES_CACERT', config.get('cacert')) or SERVICE_CERT_FILENAME
         super(Kubernetes, self).__init__({**config, 'namespace': ''}, mpp)
         if self._mpp.is_enabled():
@@ -1325,18 +1327,26 @@ class Kubernetes(AbstractDCS):
             role = None
             tmp_role = None
 
-        role_labels = {self._role_label: role}
+        updated_labels = {self._role_label: role}
         if self._tmp_role_label:
-            role_labels[self._tmp_role_label] = tmp_role
+            updated_labels[self._tmp_role_label] = tmp_role
+
+        if self._bootstrap_labels:
+            if data['state'] in ('initializing new cluster',
+                                 'running custom bootstrap script', 'starting after custom bootstrap',
+                                 'creating replica'):
+                updated_labels.update(self._bootstrap_labels)
+            else:
+                updated_labels.update({k: None for k, _ in self._bootstrap_labels.items()})
 
         member = cluster and cluster.get_member(self._name, fallback_to_leader=False)
         pod_labels = member and member.data.pop('pod_labels', None)
         ret = member and pod_labels is not None\
-            and all(pod_labels.get(k) == v for k, v in role_labels.items())\
+            and all(pod_labels.get(k) == v for k, v in updated_labels.items())\
             and deep_compare(data, member.data)
 
         if not ret:
-            metadata = {'namespace': self._namespace, 'name': self._name, 'labels': role_labels,
+            metadata: Dict[str, Any] = {'namespace': self._namespace, 'name': self._name, 'labels': updated_labels,
                                         'annotations': {'status': json.dumps(data, separators=(',', ':'))}}
             body = k8s_client.V1Pod(metadata=k8s_client.V1ObjectMeta(**metadata))
             ret = self._api.patch_namespaced_pod(self._name, self._namespace, body)

patroni/postgresql/__init__.py

@@ -713,7 +713,7 @@ class Postgresql(object):
         return time.time() - self._state_entry_timestamp
 
     def is_starting(self) -> bool:
-        return self.state == 'starting'
+        return self.state in ('starting', 'starting after custom bootstrap')
 
     def wait_for_port_open(self, postmaster: PostmasterProcess, timeout: float) -> bool:
         """Waits until PostgreSQL opens ports."""
@@ -751,9 +751,11 @@ class Postgresql(object):
         # patroni.
         self.connection_pool.close()
 
+        state = 'starting after custom bootstrap' if self.bootstrap.running_custom_bootstrap else 'starting'
+
         if self.is_running():
             logger.error('Cannot start PostgreSQL because one is already running.')
-            self.set_state('starting')
+            self.set_state(state)
             return True
 
         if not block_callbacks:
@@ -761,7 +763,7 @@ class Postgresql(object):
 
         self.set_role(role or self.get_postgres_role_from_data_directory())
 
-        self.set_state('starting')
+        self.set_state(state)
         self.set_pending_restart_reason(CaseInsensitiveDict())
 
         try:
@@ -971,7 +973,7 @@ class Postgresql(object):
     def check_startup_state_changed(self) -> bool:
         """Checks if PostgreSQL has completed starting up or failed or still starting.
 
-        Should only be called when state == 'starting'
+        Should only be called when state == 'starting [after custom bootstrap]'
 
         :returns: True if state was changed from 'starting'
         """

patroni/utils.py

@@ -925,7 +925,8 @@ def cluster_as_json(cluster: 'Cluster') -> Dict[str, Any]:
             * ``role``: ``leader``, ``standby_leader``, ``sync_standby``, ``quorum_standby``, or ``replica``;
             * ``state``: ``stopping``, ``stopped``, ``stop failed``, ``crashed``, ``running``, ``starting``,
                 ``start failed``, ``restarting``, ``restart failed``, ``initializing new cluster``, ``initdb failed``,
-                ``running custom bootstrap script``, ``custom bootstrap failed``, or ``creating replica``;
+                ``running custom bootstrap script``, ``starting after custom bootstrap``, ``custom bootstrap failed``,
+                or ``creating replica``;
             * ``api_url``: REST API URL based on ``restapi->connect_address`` configuration;
             * ``host``: PostgreSQL host based on ``postgresql->connect_address``;
             * ``port``: PostgreSQL port based on ``postgresql->connect_address``;

patroni/validator.py

@@ -1136,6 +1136,7 @@ schema = Schema({
             Optional("ports"): [{"name": str, "port": IntValidator(max=65535, expected_type=int, raise_assert=True)}],
             Optional("cacert"): str,
             Optional("retriable_http_codes"): Or(int, [int]),
+            Optional("bootstrap_labels"): dict,
         },
     }),
     Optional("citus"): {

tests/test_kubernetes.py

@@ -235,7 +235,8 @@ class BaseTestKubernetes(unittest.TestCase):
     @patch.object(k8s_client.CoreV1Api, 'list_namespaced_config_map', mock_list_namespaced_config_map, create=True)
     def setUp(self, config=None):
         config = {'ttl': 30, 'scope': 'test', 'name': 'p-0', 'loop_wait': 10, 'retry_timeout': 10,
-                  'kubernetes': {'labels': {'f': 'b'}, 'bypass_api_service': True, **(config or {})},
+                  'kubernetes': {'labels': {'f': 'b'}, 'bypass_api_service': True, **(config or {}),
+                                 'bootstrap_labels': {'foo': 'bar'}},
                   'citus': {'group': 0, 'database': 'postgres'}}
         self.k = get_dcs(config)
         self.assertIsInstance(self.k, Kubernetes)
@@ -317,10 +318,22 @@ class TestKubernetesConfigMaps(BaseTestKubernetes):
     @patch.object(k8s_client.CoreV1Api, 'patch_namespaced_pod', create=True)
     def test_touch_member(self, mock_patch_namespaced_pod):
         mock_patch_namespaced_pod.return_value.metadata.resource_version = '10'
-        self.k.touch_member({'role': 'replica'})
+
         self.k._name = 'p-1'
+        self.k.touch_member({'role': 'replica', 'state': 'initializing new cluster'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], 'bar')
+
         self.k.touch_member({'state': 'running', 'role': 'replica'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], None)
+
+        self.k.touch_member({'role': 'replica', 'state': 'running custom bootstrap script'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], 'bar')
+
+        self.k.touch_member({'role': 'replica', 'state': 'starting after custom bootstrap'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], 'bar')
+
         self.k.touch_member({'state': 'stopped', 'role': 'primary'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], None)
 
         self.k._role_label = 'isMaster'
         self.k._leader_label_value = 'true'
@@ -328,20 +341,23 @@ class TestKubernetesConfigMaps(BaseTestKubernetes):
         self.k._standby_leader_label_value = 'false'
         self.k._tmp_role_label = 'tmp_role'
 
+        self.k.touch_member({'state': 'creating replica', 'role': 'replica'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], 'bar')
+
         self.k.touch_member({'state': 'running', 'role': 'replica'})
-        mock_patch_namespaced_pod.assert_called()
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['foo'], None)
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['isMaster'], 'false')
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['tmp_role'], 'replica')
         mock_patch_namespaced_pod.rest_mock()
 
         self.k._name = 'p-0'
-        self.k.touch_member({'role': 'standby_leader'})
+        self.k.touch_member({'state': 'running', 'role': 'standby_leader'})
         mock_patch_namespaced_pod.assert_called()
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['isMaster'], 'false')
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['tmp_role'], 'primary')
         mock_patch_namespaced_pod.rest_mock()
 
-        self.k.touch_member({'role': 'primary'})
+        self.k.touch_member({'state': 'running', 'role': 'primary'})
         mock_patch_namespaced_pod.assert_called()
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['isMaster'], 'true')
         self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.labels['tmp_role'], 'primary')

tests/test_validator.py

@@ -70,6 +70,7 @@ config = {
     "kubernetes": {
         "namespace": "string",
         "labels": {},
+        'bootstrap_labels': {'foo': 'bar'},
         "scope_label": "string",
         "role_label": "string",
         "use_endpoints": False,