Convert states to enums (patroni#3293)

- Postgresql._state
- pg_isready state

Author: hughcapet

patroni/api.py

@@ -30,7 +30,7 @@
 from .__main__ import Patroni
 from .dcs import Cluster
 from .exceptions import PostgresConnectionException, PostgresException
-from .postgresql.misc import postgres_version_to_int, PostgresqlRole, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlState
 from .utils import cluster_as_json, deep_compare, enable_keepalive, parse_bool, \
     parse_int, patch_config, Retry, RetryFailedError, split_host_port, tzutc, uri
 
@@ -323,7 +323,7 @@ def do_GET(self, write_status_code_only: bool = False) -> None:
         is_lagging = leader_optime and leader_optime > replayed_location + max_replica_lag
 
         replica_status_code = 200 if not patroni.noloadbalance and not is_lagging and \
-            response.get('role') == PostgresqlRole.REPLICA and response.get('state') == PostgresqlState.RUNNING else 503
+            response.get('role') == 'replica' and response.get('state') == PostgresqlState.RUNNING else 503
 
         if not cluster and response.get('pause'):
             leader_status_code = 200 if response.get('role') in (PostgresqlRole.PRIMARY,
@@ -359,7 +359,7 @@ def do_GET(self, write_status_code_only: bool = False) -> None:
         elif 'read-only' in path and 'sync' not in path and 'quorum' not in path:
             status_code = 200 if 200 in (primary_status_code, standby_leader_status_code) else replica_status_code
         elif 'health' in path:
-            status_code = 200 if response.get('state') == 'running' else 503
+            status_code = 200 if response.get('state') == PostgresqlState.RUNNING else 503
         elif cluster:  # dcs is available
             is_quorum = response.get('quorum_standby')
             is_synchronous = response.get('sync_standby')
@@ -463,7 +463,7 @@ def do_GET_readiness(self) -> None:
         patroni = self.server.patroni
         if patroni.ha.is_leader():
             status_code = 200
-        elif patroni.postgresql.state == 'running':
+        elif patroni.postgresql.state == PostgresqlState.RUNNING:
             status_code = 200 if patroni.dcs.cluster else 503
         else:
             status_code = 503
@@ -573,7 +573,8 @@ def do_GET_metrics(self) -> None:
 
         metrics.append("# HELP patroni_postgres_running Value is 1 if Postgres is running, 0 otherwise.")
         metrics.append("# TYPE patroni_postgres_running gauge")
-        metrics.append("patroni_postgres_running{0} {1}".format(labels, int(postgres['state'] == 'running')))
+        metrics.append("patroni_postgres_running{0} {1}".format(
+            labels, int(postgres['state'] == PostgresqlState.RUNNING)))
 
         metrics.append("# HELP patroni_postmaster_start_time Epoch seconds since Postgres started.")
         metrics.append("# TYPE patroni_postmaster_start_time gauge")
@@ -899,7 +900,7 @@ def do_POST_restart(self) -> None:
             If it's not able to parse the request body, then the request is silently discarded.
         """
         status_code = 500
-        data = 'restart failed'
+        data = PostgresqlState.RESTART_FAILED
         request = self._read_json_content(body_is_optional=True)
         cluster = self.server.patroni.dcs.get_cluster()
         if request is None:
@@ -1297,10 +1298,7 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
 
         :returns: a dict with the status of Postgres/Patroni. The keys are:
 
-            * ``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``;
+            * ``state``: one of :class:`~patroni.postgresql.misc.PostgresqlState` 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;
@@ -1342,7 +1340,8 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
         config = global_config.from_cluster(cluster)
         try:
 
-            if postgresql.state not in ('running', 'restarting', 'starting'):
+            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') + ", " +\
@@ -1393,7 +1392,7 @@ def get_postgresql_status(self, retry: bool = False) -> Dict[str, Any]:
 
         except (psycopg.Error, RetryFailedError, PostgresConnectionException):
             state = postgresql.state
-            if state == 'running':
+            if state == PostgresqlState.RUNNING:
                 logger.exception('get_postgresql_status')
                 state = 'unknown'
             result: Dict[str, Any] = {'state': state, 'role': postgresql.role}

patroni/ctl.py

@@ -65,7 +65,7 @@
 from .config import Config
 from .dcs import AbstractDCS, Cluster, get_dcs as _get_dcs, Member
 from .exceptions import PatroniException
-from .postgresql.misc import postgres_version_to_int, PostgresqlRole, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlState
 from .postgresql.mpp import get_mpp
 from .request import PatroniRequest
 from .utils import cluster_as_json, patch_config, polling_loop
@@ -1236,7 +1236,7 @@ def reinit(cluster_name: str, group: Optional[int], member_names: List[str], for
         time.sleep(2)
         for member in wait_on_members:
             data = json.loads(request_patroni(member, 'get', 'patroni').data.decode('utf-8'))
-            if data.get('state') != 'creating replica':
+            if data.get('state') != PostgresqlState.CREATING_REPLICA:
                 click.echo('Reinitialize is completed on: {0}'.format(member.name))
                 wait_on_members.remove(member)
 
@@ -1746,10 +1746,7 @@ def output_members(cluster: Cluster, name: str, extended: bool = False,
         * ``Member``: name of the Patroni node, as per ``name`` configuration;
         * ``Host``: hostname (or IP) and port, as per ``postgresql.listen`` configuration;
         * ``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;
+        * ``State``: one of :class:`~patroni.postgresql.misc.PostgresqlState`;
         * ``TL``: current timeline in Postgres;
           ``Lag in MB``: replication lag.
 

patroni/dcs/__init__.py

@@ -296,8 +296,10 @@ def state(self) -> str:
 
     @property
     def is_running(self) -> bool:
-        """``True`` if the member :attr:`~Member.state` is ``running``."""
-        return self.state == 'running'
+        """``True`` if the member :attr:`~Member.state` is :class:`~patroni.postgresql.misc.PostgresqlState.RUNNING`."""
+        from ..postgresql.misc import PostgresqlState
+
+        return self.state == PostgresqlState.RUNNING
 
     @property
     def patroni_version(self) -> Optional[Tuple[int, ...]]:

patroni/dcs/consul.py

@@ -19,7 +19,7 @@
 from urllib3.exceptions import HTTPError
 
 from ..exceptions import DCSError
-from ..postgresql.misc import PostgresqlRole, PostgresqlState
+from ..postgresql.misc import PostgresqlState
 from ..postgresql.mpp import AbstractMPP
 from ..utils import deep_compare, parse_bool, Retry, RetryFailedError, split_host_port, uri, USER_AGENT
 from . import AbstractDCS, catch_return_false_exception, Cluster, ClusterConfig, \
@@ -547,12 +547,12 @@ def _update_service(self, data: Dict[str, Any]) -> Optional[bool]:
             'enable_tag_override': True,
         }
 
-        if state == 'stopped' or (not self._register_service and self._previous_loop_register_service):
+        if state == PostgresqlState.STOPPED or (not self._register_service and self._previous_loop_register_service):
             self._previous_loop_register_service = self._register_service
             return self.deregister_service(params['service_id'])
 
         self._previous_loop_register_service = self._register_service
-        if data['role'] in [PostgresqlRole.PRIMARY, PostgresqlRole.REPLICA, PostgresqlRole.STANDBY_LEADER]:
+        if role in ['primary', 'replica', 'standby-leader']:
             if state != PostgresqlState.RUNNING:
                 return
             return self.register_service(service_name, **params)

patroni/dcs/kubernetes.py

@@ -23,7 +23,7 @@
 
 from ..collections import EMPTY_DICT
 from ..exceptions import DCSError
-from ..postgresql.misc import PostgresqlRole, PostgresqlState
+from ..postgresql.misc import PostgresqlState
 from ..postgresql.mpp import AbstractMPP
 from ..utils import deep_compare, iter_response_objects, \
     keepalive_socket_options, Retry, RetryFailedError, tzutc, uri, USER_AGENT
@@ -1324,7 +1324,7 @@ def touch_member(self, data: Dict[str, Any]) -> bool:
             role = self._standby_leader_label_value \
                 if data['role'] == PostgresqlRole.STANDBY_LEADER else self._leader_label_value
             tmp_role = 'primary'
-        elif data['state'] == PostgresqlState.RUNNING and data['role'] != PostgresqlRole.PRIMARY:
+        elif data['state'] == PostgresqlState.RUNNING and data['role'] != 'primary':
             role = {'replica': self._follower_label_value}.get(data['role'], data['role'])
             tmp_role = data['role']
         else:
@@ -1333,7 +1333,14 @@ def touch_member(self, data: Dict[str, Any]) -> bool:
 
         role_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 (PostgresqlState.INITDB, PostgresqlState.CUSTOM_BOOTSTRAP,
+                                 PostgresqlState.BOOTSTRAP_STARTING, PostgresqlState.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)

patroni/ha.py

@@ -17,7 +17,7 @@
 from .dcs import AbstractDCS, Cluster, Leader, Member, RemoteMember, Status, SyncState
 from .exceptions import DCSError, PatroniFatalException, PostgresConnectionException
 from .postgresql.callback_executor import CallbackAction
-from .postgresql.misc import postgres_version_to_int, PostgresqlRole, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlState
 from .postgresql.postmaster import PostmasterProcess
 from .postgresql.rewind import Rewind
 from .quorum import QuorumStateResolver
@@ -474,7 +474,8 @@ def touch_member(self) -> bool:
                 data['pending_restart'] = True
                 data['pending_restart_reason'] = dict(self.state_handler.pending_restart_reason)
             if self._async_executor.scheduled_action in (None, 'promote') \
-                    and data['state'] in ['running', 'restarting', 'starting']:
+                    and data['state'] in [PostgresqlState.RUNNING, PostgresqlState.RESTARTING,
+                                          PostgresqlState.STARTING]:
                 try:
                     timeline, wal_position, pg_control_timeline = self.state_handler.timeline_wal_position()
                     data['xlog_location'] = self._last_wal_lsn = wal_position
@@ -514,7 +515,7 @@ def touch_member(self) -> bool:
             ret = self.dcs.touch_member(data)
             if ret:
                 new_state = (data['state'], data['role'])
-                if self._last_state != new_state and new_state == (PostgresqlState.RUNNING, PostgresqlRole.PRIMARY):
+                if self._last_state != new_state and new_state == (PostgresqlState.RUNNING, 'primary'):
                     self.notify_mpp_coordinator('after_promote')
                 self._last_state = new_state
             return ret
@@ -660,7 +661,7 @@ def recover(self) -> str:
                 and data.get('Database cluster state') in ('in production', 'in crash recovery',
                                                            'shutting down', 'shut down')\
                 and self.state_handler.state == PostgresqlState.CRASHED\
-                and self.state_handler.role == PostgresqlRole.PRIMARY\
+                and self.state_handler.role == 'primary'\
                 and not self.state_handler.config.recovery_conf_exists():
             # We know 100% that we were running as a primary a few moments ago, therefore could just start postgres
             msg = 'starting primary after failure'
@@ -1183,7 +1184,7 @@ def update_failsafe(self, data: Dict[str, Any]) -> Union[int, str, None]:
 
         :returns: the reason why caller shouldn't continue as a primary or the current value of received/replayed LSN.
         """
-        if self.state_handler.state == PostgresqlState.RUNNING and self.state_handler.role == PostgresqlRole.PRIMARY:
+        if self.state_handler.state == PostgresqlState.RUNNING and self.state_handler.role == 'primary':
             return 'Running as a leader'
         self._failsafe.update(data)
         return self._last_wal_lsn
@@ -1919,7 +1920,7 @@ def after_start() -> None:
             elif res is None:
                 return (False, 'postgres is still starting')
             else:
-                return (False, 'restart failed')
+                return (False, PostgresqlState.RESTART_FAILED)
 
     def _do_reinitialize(self, cluster: Cluster) -> Optional[bool]:
         self.state_handler.stop('immediate', stop_timeout=self.patroni.config['retry_timeout'])
@@ -2059,7 +2060,8 @@ def handle_starting_instance(self) -> Optional[str]:
         if not self.state_handler.check_for_startup() or self.is_paused():
             self.set_start_timeout(None)
             if self.is_paused():
-                self.state_handler.set_state(self.state_handler.is_running() and 'running' or 'stopped')
+                self.state_handler.set_state(PostgresqlState.RUNNING if self.state_handler.is_running()
+                                             else PostgresqlState.STOPPED)
             return None
 
         # state_handler.state == 'starting' here
@@ -2225,7 +2227,7 @@ def _run_cycle(self) -> str:
 
             if not self.state_handler.is_healthy():
                 if self.is_paused():
-                    self.state_handler.set_state('stopped')
+                    self.state_handler.set_state(PostgresqlState.STOPPED)
                     if self.has_lock():
                         self._delete_leader()
                         return 'removed leader lock because postgres is not running'
@@ -2238,8 +2240,8 @@ def _run_cycle(self) -> str:
                             (self._rewind.is_needed and self._rewind.can_rewind_or_reinitialize_allowed):
                         return 'postgres is not running'
 
-                if self.state_handler.state in ('running', 'starting'):
-                    self.state_handler.set_state('crashed')
+                if self.state_handler.state in (PostgresqlState.RUNNING, PostgresqlState.STARTING):
+                    self.state_handler.set_state(PostgresqlState.CRASHED)
                 # try to start dead postgres
                 return self.recover()