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

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}'