Convert roles to enums (patroni#3303)

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, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlRole, 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') == 'replica' and response.get('state') == PostgresqlState.RUNNING else 503
+            response.get('role') == PostgresqlRole.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,

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, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlRole, PostgresqlState
 from .postgresql.mpp import get_mpp
 from .request import PatroniRequest
 from .utils import cluster_as_json, patch_config, polling_loop

patroni/dcs/consul.py

@@ -19,7 +19,7 @@
 from urllib3.exceptions import HTTPError
 
 from ..exceptions import DCSError
-from ..postgresql.misc import PostgresqlState
+from ..postgresql.misc import PostgresqlRole, 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, \
@@ -552,7 +552,7 @@ def _update_service(self, data: Dict[str, Any]) -> Optional[bool]:
             return self.deregister_service(params['service_id'])
 
         self._previous_loop_register_service = self._register_service
-        if role in ['primary', 'replica', 'standby-leader']:
+        if data['role'] in [PostgresqlRole.PRIMARY, PostgresqlRole.REPLICA, PostgresqlRole.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 PostgresqlState
+from ..postgresql.misc import PostgresqlRole, 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'] != 'primary':
+        elif data['state'] == PostgresqlState.RUNNING and data['role'] != PostgresqlRole.PRIMARY:
             role = {'replica': self._follower_label_value}.get(data['role'], data['role'])
             tmp_role = data['role']
         else:

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, PostgresqlState
+from .postgresql.misc import postgres_version_to_int, PostgresqlRole, PostgresqlState
 from .postgresql.postmaster import PostmasterProcess
 from .postgresql.rewind import Rewind
 from .quorum import QuorumStateResolver
@@ -515,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, 'primary'):
+                if self._last_state != new_state and new_state == (PostgresqlState.RUNNING, PostgresqlRole.PRIMARY):
                     self.notify_mpp_coordinator('after_promote')
                 self._last_state = new_state
             return ret
@@ -661,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 == 'primary'\
+                and self.state_handler.role == PostgresqlRole.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'
@@ -1184,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 == 'primary':
+        if self.state_handler.state == PostgresqlState.RUNNING and self.state_handler.role == PostgresqlRole.PRIMARY:
             return 'Running as a leader'
         self._failsafe.update(data)
         return self._last_wal_lsn

patroni/postgresql/__init__.py

@@ -28,7 +28,7 @@
 from .cancellable import CancellableSubprocess
 from .config import ConfigHandler, mtime
 from .connection import ConnectionPool, get_connection_cursor
-from .misc import parse_history, parse_lsn, postgres_major_version_to_int, PostgresqlState
+from .misc import parse_history, parse_lsn, postgres_major_version_to_int, PostgresqlRole, PostgresqlState
 from .mpp import AbstractMPP
 from .postmaster import PostmasterProcess
 from .slots import SlotsHandler

patroni/postgresql/config.py

@@ -20,7 +20,7 @@
 from ..utils import compare_values, get_postgres_version, is_subpath, \
     maybe_convert_from_base_unit, parse_bool, parse_int, split_host_port, uri, validate_directory
 from ..validator import EnumValidator, IntValidator
-from .misc import get_major_from_minor_version, postgres_version_to_int, PostgresqlState
+from .misc import get_major_from_minor_version, postgres_version_to_int, PostgresqlRole, PostgresqlState
 from .validator import recovery_parameters, transform_postgresql_parameter_value, transform_recovery_parameter_value
 
 if TYPE_CHECKING:  # pragma: no cover

patroni/postgresql/misc.py

@@ -34,6 +34,22 @@ def __repr__(self) -> str:
         return self.value
 
 
+class PostgresqlRole(str, Enum):
+    """Possible values of :attr:`Postgresql.role`."""
+
+    PRIMARY = 'primary'
+    MASTER = 'master'
+    STANDBY_LEADER = 'standby_leader'
+    REPLICA = 'replica'
+    DEMOTED = 'demoted'
+    UNINITIALIZED = 'uninitialized'
+    PROMOTED = 'promoted'
+
+    def __repr__(self) -> str:
+        """Get a string representation of a :class:`PostgresqlRole` member."""
+        return self.value
+
+
 def postgres_version_to_int(pg_version: str) -> int:
     """Convert the server_version to integer
 

patroni/postgresql/mpp/citus.py

@@ -9,7 +9,7 @@
 from ...dcs import Cluster
 from ...psycopg import connect, ProgrammingError, quote_ident
 from ...utils import parse_int
-from ..misc import PostgresqlState
+from ..misc import PostgresqlRole, PostgresqlState
 from . import AbstractMPP, AbstractMPPHandler
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -478,7 +478,7 @@ def sync_meta_data(self, cluster: Cluster) -> None:
         for groupid, worker in cluster.workers.items():
             leader = worker.leader
             if leader and leader.conn_url\
-                    and leader.data.get('role') in ('master', 'primary')\
+                    and leader.data.get('role') in (PostgresqlRole.MASTER, PostgresqlRole.PRIMARY)\
                     and leader.data.get('state') == PostgresqlState.RUNNING:
                 self.add_task('after_promote', groupid, worker, leader.name, leader.conn_url)
 

tests/test_api.py

@@ -15,7 +15,7 @@
 from patroni.exceptions import PostgresConnectionException
 from patroni.ha import _MemberStatus
 from patroni.postgresql.config import get_param_diff
-from patroni.postgresql.misc import PostgresqlState
+from patroni.postgresql.misc import PostgresqlRole, PostgresqlState
 from patroni.psycopg import OperationalError
 from patroni.utils import RetryFailedError, tzutc
 
@@ -51,7 +51,7 @@ class MockPostgresql:
     connection_pool = MockConnectionPool()
     name = 'test'
     state = PostgresqlState.RUNNING
-    role = 'primary'
+    role = PostgresqlRole.PRIMARY
     server_version = 90625
     major_version = 90600
     sysid = 'dummysysid'