wip

patroni/config.py

@@ -657,7 +657,7 @@ def _get_auth(name: str, params: Collection[str] = _AUTH_ALLOWED_PARAMETERS[:2])
                               '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_ANNOTATIONS') and name:
                     value = os.environ.pop(param)
                     if name == 'CITUS':
                         if suffix == 'GROUP':
@@ -668,7 +668,7 @@ def _get_auth(name: str, params: Collection[str] = _AUTH_ALLOWED_PARAMETERS[:2])
                         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_ANNOTATIONS'):
                         value = _parse_dict(value)
                     elif suffix in ('USE_PROXIES', 'REGISTER_SERVICE', 'USE_ENDPOINTS', 'BYPASS_API_SERVICE', 'VERIFY'):
                         value = parse_bool(value)

patroni/dcs/kubernetes.py

@@ -760,6 +760,7 @@ def __init__(self, config: Dict[str, Any], mpp: AbstractMPP) -> None:
         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_annotations: Dict[str, str] = {str(k): str(v) for k, v in (config.get('bootstrap_annotations') 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():
@@ -1336,8 +1337,13 @@ def touch_member(self, data: Dict[str, Any]) -> bool:
             and deep_compare(data, member.data)
 
         if not ret:
-            metadata = {'namespace': self._namespace, 'name': self._name, 'labels': role_labels,
-                        'annotations': {'status': json.dumps(data, separators=(',', ':'))}}
+            metadata: Dict[str, Any] = {'namespace': self._namespace, 'name': self._name, 'labels': role_labels,
+                                        'annotations': {'status': json.dumps(data, separators=(',', ':'))}}
+            if self._bootstrap_annotations:
+                if data['state'] in ('initializing new cluster', 'running custom bootstrap script', 'creating replica'):
+                    metadata['annotations'].update(self._bootstrap_annotations)
+                else:
+                    metadata['annotations'].update({k: None for k, _ in self._bootstrap_annotations.items()})
             body = k8s_client.V1Pod(metadata=k8s_client.V1ObjectMeta(**metadata))
             ret = self._api.patch_namespaced_pod(self._name, self._namespace, body)
             if ret:

patroni/postgresql/bootstrap.py

@@ -115,6 +115,7 @@ def option_is_allowed(name: str) -> bool:
 
     def _initdb(self, config: Any) -> bool:
         self._postgresql.set_state('initializing new cluster')
+        time.sleep(15)
         not_allowed_options = ('pgdata', 'nosync', 'pwfile', 'sync-only', 'version')
 
         def error_handler(e: str) -> None:

patroni/validator.py

@@ -1135,6 +1135,7 @@ def validate_watchdog_mode(value: Any) -> None:
             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_annotations"): 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_annotations': {'foo': 'bar'}},
                   'citus': {'group': 0, 'database': 'postgres'}}
         self.k = get_dcs(config)
         self.assertIsInstance(self.k, Kubernetes)
@@ -317,31 +318,43 @@ def test_set_config_value(self):
     @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.annotations['foo'], 'bar')
+
         self.k.touch_member({'state': 'running', 'role': 'replica'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.annotations['foo'], None)
+
+        self.k.touch_member({'role': 'replica', 'state': 'running custom bootstrap script'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.annotations['foo'], 'bar')
+
         self.k.touch_member({'state': 'stopped', 'role': 'primary'})
+        self.assertEqual(mock_patch_namespaced_pod.call_args[0][2].metadata.annotations['foo'], None)
 
         self.k._role_label = 'isMaster'
         self.k._leader_label_value = 'true'
         self.k._follower_label_value = 'false'
         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.annotations['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.annotations['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 @@
     "kubernetes": {
         "namespace": "string",
         "labels": {},
+        'bootstrap_annotations': {'foo': 'bar'},
         "scope_label": "string",
         "role_label": "string",
         "use_endpoints": False,