Merge branch 'release/v3.3.0' of github.com:zalando/patroni into feature/citus-secondaries

Most of them are related to partially unknown types of values from empty
dict or list. To solve it for the empty dict we use `EMPTY_DICT` object of
newly introduced `_FrozenDict` class.

docs/releases.rst

@@ -6,16 +6,6 @@ Release notes
 Version 3.3.0
 -------------
 
-.. warning::
-   All older Partoni versions are not compatible with ``ydiff>=1.3``.
-
-   There are the following options available to "fix" the problem:
-
-   1. upgrade Patroni to the latest version
-   2. install ``ydiff<1.3`` after installing Patroni
-   3. install ``cdiff`` module
-
-
 **New features**
 
 - Add ability to pass ``auth_data`` to Zookeeper client (Aras Mumcuyan)
@@ -26,9 +16,9 @@ Version 3.3.0
 
   Provide an application ``patroni_barman`` that allows to perform ``Barman`` operations remotely and can be used as a custom bootstrap/custom replica method or as an ``on_role_change`` callback. Please check :ref:`here <tools_integration>` for more information.
 
-- Support ``JSON`` log format (alisalemmi)
+- Support JSON log format (alisalemmi)
 
-  Apart from ``plain`` (default), Patroni now also supports ``json`` log format. Requires ``python-json-logger>=2.0.2`` library to be installed.
+  Apart from ``plain``, Patroni now also supports ``json`` log format. Requires ``python-json-logger`` library to be installed.
 
 - Show ``pending_restart_reason`` information (Polina Bungina)
 
@@ -41,7 +31,7 @@ Version 3.3.0
 
 **Improvements**
 
-- Implement validation of the ``log`` section (Alexander Kukushkin)
+- Implement validation of the log section (Alexander Kukushkin)
 
   Until now validator was not checking the correctness of the logging configuration provided.
 
@@ -70,7 +60,8 @@ Version 3.3.0
 
 - Add compatibility code for ``ydiff`` library version 1.3+ (Alexander Kukushkin)
 
-  Patroni is relying on some API from ``ydiff`` that is not public because it is supposed to be just a terminal tool rather than a python module. Unfortunately, the API change in 1.3 broke old Patroni versions.
+.. warning::
+    All older Partoni versions are not compatible with ``ydiff`` 1.3+. Please upgrade Patroni, use ``ydiff`` version <1.3, or install ``cdiff``.
 
 
 Version 3.2.2

features/citus.feature

@@ -11,7 +11,9 @@ Feature: citus
     Then replication works from postgres0 to postgres1 after 15 seconds
     Then replication works from postgres2 to postgres3 after 15 seconds
     And postgres0 is registered in the postgres0 as the primary in group 0 after 5 seconds
+    And postgres1 is registered in the postgres0 as the secondary in group 0 after 5 seconds
     And postgres2 is registered in the postgres0 as the primary in group 1 after 5 seconds
+    And postgres3 is registered in the postgres0 as the secondary in group 1 after 5 seconds
 
   Scenario: coordinator failover updates pg_dist_node
     Given I run patronictl.py failover batman --group 0 --candidate postgres1 --force
@@ -19,11 +21,13 @@ Feature: citus
     And "members/postgres0" key in a group 0 in DCS has state=running after 15 seconds
     And replication works from postgres1 to postgres0 after 15 seconds
     And postgres1 is registered in the postgres2 as the primary in group 0 after 5 seconds
+    And postgres0 is registered in the postgres2 as the secondary in group 0 after 15 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres0 after 15 seconds
     When I run patronictl.py switchover batman --group 0 --candidate postgres0 --force
     Then postgres0 role is the primary after 10 seconds
     And replication works from postgres0 to postgres1 after 15 seconds
     And postgres0 is registered in the postgres2 as the primary in group 0 after 5 seconds
+    And postgres1 is registered in the postgres2 as the secondary in group 0 after 15 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres1 after 15 seconds
 
   Scenario: worker switchover doesn't break client queries on the coordinator
@@ -35,6 +39,7 @@ Feature: citus
     And "members/postgres2" key in a group 1 in DCS has state=running after 15 seconds
     And replication works from postgres3 to postgres2 after 15 seconds
     And postgres3 is registered in the postgres0 as the primary in group 1 after 5 seconds
+    And postgres2 is registered in the postgres0 as the secondary in group 1 after 15 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres2 after 15 seconds
     And a thread is still alive
     When I run patronictl.py switchover batman --group 1 --force
@@ -42,6 +47,7 @@ Feature: citus
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
     And postgres2 is registered in the postgres0 as the primary in group 1 after 5 seconds
+    And postgres3 is registered in the postgres0 as the secondary in group 1 after 15 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres3 after 15 seconds
     And a thread is still alive
     When I stop a thread
@@ -55,6 +61,7 @@ Feature: citus
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
     And postgres2 is registered in the postgres0 as the primary in group 1 after 5 seconds
+    And postgres3 is registered in the postgres0 as the secondary in group 1 after 15 seconds
     And a thread is still alive
     When I stop a thread
     Then a distributed table on postgres0 has expected rows

patroni.spec

@@ -13,17 +13,13 @@ def hiddenimports():
         sys.path.pop(0)
 
 
-def resources():
-    import os
-    res_dir = 'patroni/postgresql/available_parameters/'
-    exts = set(f.split('.')[-1] for f in os.listdir(res_dir))
-    return [(res_dir + '*.' + e, res_dir) for e in exts if e.lower() in {'yml', 'yaml'}]
-
-
 a = Analysis(['patroni/__main__.py'],
              pathex=[],
              binaries=None,
-             datas=resources(),
+             datas=[
+                ('patroni/postgresql/available_parameters/*.yml', 'patroni/postgresql/available_parameters'),
+                ('patroni/postgresql/available_parameters/*.yaml', 'patroni/postgresql/available_parameters'),
+            ],
              hiddenimports=hiddenimports(),
              hookspath=[],
              runtime_hooks=[],

patroni/postgresql/mpp/citus.py

@@ -4,7 +4,7 @@ import time
 
 from threading import Condition, Event, Thread
 from urllib.parse import urlparse
-from typing import Any, Dict, List, Optional, Union, Tuple, TYPE_CHECKING
+from typing import Any, Collection, Dict, Iterator, List, Optional, Union, Set, Tuple, TYPE_CHECKING
 
 from . import AbstractMPP, AbstractMPPHandler
 from ...dcs import Cluster
@@ -19,19 +19,312 @@ CITUS_SLOT_NAME_RE = re.compile(r'^citus_shard_(move|split)_slot(_[1-9][0-9]*){2
 logger = logging.getLogger(__name__)
 
 
-class PgDistNode(object):
-    """Represents a single row in the `pg_dist_node` table"""
+class PgDistNode:
+    """Represents a single row in "pg_dist_node" table.
 
-    def __init__(self, group: int, host: str, port: int, event: str, nodeid: Optional[int] = None,
-                 timeout: Optional[float] = None, cooldown: Optional[float] = None) -> None:
-        self.group = group
-        # A weird way of pausing client connections by adding the `-demoted` suffix to the hostname
-        self.host = host + ('-demoted' if event == 'before_demote' else '')
+    .. note::
+
+        Unlike "noderole" possible values of ``role`` are ``primary``, ``secondary``, and ``demoted``.
+        The last one is used to pause client connections on the coordinator to the worker by
+        appending ``-demoted`` suffix to the "nodename". The actual "noderole" in DB remains ``primary``.
+
+    :ivar host: "nodename" value
+    :ivar port: "nodeport" value
+    :ivar role: "noderole" value
+    :ivar nodeid: "nodeid" value
+    """
+
+    def __init__(self, host: str, port: int, role: str, nodeid: Optional[int] = None) -> None:
+        """Create a :class:`PgDistNode` object based on given arguments.
+
+        :param host: "nodename" of the Citus coordinator or worker.
+        :param port: "nodeport" of the Citus coordinator or worker.
+        :param role: "noderole" value.
+        :param nodeid: id of the row in the "pg_dist_node".
+        """
+        self.host = host
         self.port = port
+        self.role = role
+        self.nodeid = nodeid
+
+    def __hash__(self) -> int:
+        """Defines a hash function to put :class:`PgDistNode` objects to :class:`PgDistGroup` set-like object.
+
+        .. note::
+            We use (:attr:`host`, :attr:`port`) tuple here because it is one of the UNIQUE constraints on the
+            "pg_dist_node" table. The :attr:`role` value is irrelevant here because nodes may change their roles.
+        """
+        return hash((self.host, self.port))
+
+    def __eq__(self, other: Any) -> bool:
+        """Defines a comparison function.
+
+        :returns: ``True`` if :attr:`host` and :attr:`port` between two instances are the same.
+        """
+        return isinstance(other, PgDistNode) and self.host == other.host and self.port == other.port
+
+    def __str__(self) -> str:
+        return ('PgDistNode(nodeid={0},host={1},port={2},role={3})'
+                .format(self.nodeid, self.host, self.port, self.role))
+
+    def __repr__(self) -> str:
+        return str(self)
+
+    def is_primary(self) -> bool:
+        """Checks whether this object represents "primary" in a corresponding group.
+
+        :returns: ``True`` if this object represents the ``primary``.
+        """
+        return self.role in ('primary', 'demoted')
+
+    def as_tuple(self, include_nodeid: bool = False) -> Tuple[str, int, str, Optional[int]]:
+        """Helper method to compare two :class:`PgDistGroup` objects.
+
+        .. note::
+
+            *include_nodeid* is set to ``True`` only in unit-tests.
+
+        :param include_nodeid: whether :attr:`nodeid` should be taken into account when comparison is performed.
+
+        :returns: :class:`tuple` object with :attr:`host`, :attr:`port`, :attr:`role`, and optionally :attr:`nodeid`.
+        """
+        return self.host, self.port, self.role, (self.nodeid if include_nodeid else None)
+
+
+class PgDistGroup(Set[PgDistNode]):
+    """A :class:`set`-like object that represents a Citus group in "pg_dist_node" table.
+
+    This class implements a set of methods to compare topology and if it is necessary
+    to transition from the old to the new topology in a "safe" manner:
+
+        * register new primary/secondaries
+        * replace gone secondaries with added secondaries
+        * failover and switchover
+
+    Typically there will be at least one :class:`PgDistNode` object registered (``primary``).
+    In addition to that there could be one or more ``secondary`` nodes.
+
+    :ivar failover: whether the ``primary`` row should be updated as a result of :func:`transition` method call.
+    :ivar groupid: the "groupid" from "pg_dist_node".
+    """
+
+    def __init__(self, groupid: int, nodes: Optional[Collection[PgDistNode]] = None) -> None:
+        """Creates a :class:`PgDistGroup` object based on given arguments.
+
+        :param groupid: the groupid from "pg_dist_node".
+        :param nodes: a collection of :class:`PgDistNode` objects that belog to a *groupid*.
+        """
+        self.failover = False
+        self.groupid = groupid
+
+        if nodes:
+            self.update(nodes)
+
+    def equals(self, other: 'PgDistGroup', check_nodeid: bool = False) -> bool:
+        """Compares two :class:`PgDistGroup` objects.
+
+        :param other: what we want to compare with.
+        :param check_nodeid: whether :attr:`PgDistNode.nodeid` should be compared in addition to
+                             :attr:`PgDistNode.host`, :attr:`PgDistNode.port`, and :attr:`PgDistNode.role`.
+
+        :returns: ``True`` if two :class:`PgDistGroup` objects are fully identical.
+        """
+        return self.groupid == other.groupid\
+            and set(v.as_tuple(check_nodeid) for v in self) == set(v.as_tuple(check_nodeid) for v in other)
+
+    def primary(self) -> Optional[PgDistNode]:
+        """Finds and returns :class:`PgDistNode` object that represents the "primary".
+
+        :returns: :class:`PgDistNode` object which represents the "primary" or ``None`` if not found.
+        """
+        return next(iter(v for v in self if v.is_primary()), None)
+
+    def get(self, value: PgDistNode) -> Optional[PgDistNode]:
+        """Performs a lookup of the actual value in a set.
+
+        .. note::
+            It is necessary because :func:`__hash__` and :func:`__eq__` methods in :class:`PgDistNode` are
+            redefined and effectively they check only :attr:`PgDistNode.host` and :attr:`PgDistNode.port` attributes.
+
+        :param value: the key we search for.
+        :returns: the actual :class:`PgDistNode` value from this :class:`PgDistGroup` object or ``None`` if not found.
+        """
+        return next(iter(v for v in self if v == value), None)
+
+    def transition(self, old: 'PgDistGroup') -> Iterator[PgDistNode]:
+        """Compares this topology with the old one and yields transitions that transform the old to the new one.
+
+        .. note::
+            The actual yielded object is :class:`PgDistNode` that will be passed to
+            the :meth:`CitusHandler.update_node` to execute all transitions in a transaction.
+
+            In addition to the yielding transactions this method fills up :attr:`PgDistNode.nodeid`
+            attribute for nodes that are presented in the old and in the new topology.
+
+            There are a few simple rules/constraints that are imposed by Citus and must be followed:
+            - adding/removing nodes is only possible when metadata is synced to all registered "priorities".
+
+            - the "primary" row in "pg_dist_node" always keeps the nodeid (unless it is
+              removed, but it is not supported by Patroni).
+
+            - "nodename", "nodeport" must be unique across all rows in the "pg_dist_node". This means that
+              every time we want to change the nodeid of an existing node (i.e. to change it from secondary
+              to primary), we should first write some other "nodename"/"nodeport" to the row it's currently in.
+
+            - updating "broken" nodes always works and metadata is synced asynchnonously after the commit.
+
+        Following these rules below is an example of the switchover between node1 (primary, nodeid=4)
+        and node2 (secondary, nodeid=5).
+
+        .. code-block:: SQL
+
+            BEGIN;
+                SELECT citus_update_node(4, 'node1-demoted', 5432);
+                SELECT citus_update_node(5, 'node1', 5432);
+                SELECT citus_update_node(4, 'node2', 5432);
+            COMMIT;
+
+        :param old: the last known topology registered in "pg_dist_node" for a given :attr:`groupid`.
+
+        :yields: :class:`PgDistNode` objects that must be updated/added/removed in "pg_dist_node".
+        """
+        self.failover = old.failover
+
+        new_primary = self.primary()
+        assert new_primary is not None
+        old_primary = old.primary()
+
+        gone_nodes = old - self - {old_primary}
+        added_nodes = self - old - {new_primary}
+
+        if not old_primary:
+            # We did not have any nodes in the group yet and we're adding one now
+            yield new_primary
+        elif old_primary == new_primary:
+            new_primary.nodeid = old_primary.nodeid
+            # Controlled switchover with pausing client connections.
+            # Achieved by updating the primary row and putting hostname = '${host}-demoted' in a transaction.
+            if old_primary.role != new_primary.role:
+                self.failover = True
+                yield new_primary
+        elif old_primary != new_primary:
+            self.failover = True
+
+            new_primary_old_node = old.get(new_primary)
+            old_primary_new_node = self.get(old_primary)
+
+            # The new primary was registered as a secondary before failover
+            if new_primary_old_node:
+                new_node = None
+                # Old primary is gone and some new secondaries were added.
+                # We can use the row of promoted secondary to add the new secondary.
+                if not old_primary_new_node and added_nodes:
+                    new_node = added_nodes.pop()
+                    new_node.nodeid = new_primary_old_node.nodeid
+                    yield new_node
+
+                    # notify _maybe_register_old_primary_as_secondary that the old primary should not be re-registered
+                    old_primary.role = 'secondary'
+                # In opposite case we need to change the primary record to '${host}-demoted:${port}'
+                # before we can put its host:port to the row of promoted secondary.
+                elif old_primary.role == 'primary':
+                    old_primary.role = 'demoted'
+                    yield old_primary
+
+                # The old primary is gone and the promoted secondary row wasn't yet used.
+                if not old_primary_new_node and not new_node:
+                    # We have to "add" the gone primary to the row of promoted secondary because
+                    # nodes could not be removed while the metadata isn't synced.
+                    old_primary_new_node = PgDistNode(old_primary.host, old_primary.port, new_primary_old_node.role)
+                    self.add(old_primary_new_node)
+
+                # put the old primary instead of promoted secondary
+                if old_primary_new_node:
+                    old_primary_new_node.nodeid = new_primary_old_node.nodeid
+                    yield old_primary_new_node
+
+            # update the primary record with the new information
+            new_primary.nodeid = old_primary.nodeid
+            yield new_primary
+
+            # The new primary was never registered as a standby and there are secondaries that have gone away. Since
+            # nodes can't be removed while metadata isn't synced we have to temporarily "add" the old primary back.
+            if not new_primary_old_node and gone_nodes:
+                # We were in the middle of controlled switchover while the primary disappeared.
+                # If there are any gone nodes that can't be reused for new secondaries we will
+                # use one of them to temporarily "add" the old primary back as a secondary.
+                if not old_primary_new_node and old_primary.role == 'demoted' and len(gone_nodes) > len(added_nodes):
+                    old_primary_new_node = PgDistNode(old_primary.host, old_primary.port, 'secondary')
+                    self.add(old_primary_new_node)
+
+                # Use one of the gone secondaries to put host:port of the old primary there.
+                if old_primary_new_node:
+                    old_primary_new_node.nodeid = gone_nodes.pop().nodeid
+                    yield old_primary_new_node
+
+        # Fill nodeid for standbys in the new topology from the old ones
+        old_replicas = {v: v for v in old if not v.is_primary()}
+        for n in self:
+            if not n.is_primary() and not n.nodeid and n in old_replicas:
+                n.nodeid = old_replicas[n].nodeid
+
+        # Reuse nodeid's of gone standbys to "add" new standbys
+        while gone_nodes and added_nodes:
+            a = added_nodes.pop()
+            a.nodeid = gone_nodes.pop().nodeid
+            yield a
+
+        # Adding or removing nodes operations are executed on primaries in all Citus groups in 2PC.
+        # If we know that the primary was updated (self.failover is True) that automatically means that
+        # adding/removing nodes calls will fail and the whole transaction will be aborted. Therefore
+        # we discard operations that add/remove secondaries if we know that the primary was just updated.
+        # The inconsistency will be automatically resolved on the next Patroni heartbeat loop.
+
+        # Remove remaining nodes that are gone, but only in case if metadata is in sync (self.failover is False).
+        for g in gone_nodes:
+            if not self.failover:
+                # Remove the node if we expect metadata to be in sync
+                yield PgDistNode(g.host, g.port, '')
+            else:
+                # Otherwise add these nodes to the new topology
+                self.add(g)
+
+        # Add new nodes to the metadata, but only in case if metadata is in sync (self.failover is False).
+        for a in added_nodes:
+            if not self.failover:
+                # Add the node if we expect metadata to be in sync
+                yield a
+            else:
+                # Otherwise remove them from the new topology
+                self.discard(a)
+
+
+class PgDistTask(PgDistGroup):
+    """A "task" that represents the current or desired state of "pg_dist_node" for a provided *groupid*.
+
+    :ivar group: the "groupid" in "pg_dist_node".
+    :ivar event: an "event" that resulted in creating this task.
+                 possible values: "before_demote", "before_promote", "after_promote".
+    :ivar timeout: a transaction timeout if the task resulted in starting a transaction.
+    :ivar cooldown: the cooldown value for ``citus_update_node()`` UDF call.
+    :ivar deadline: the time in unix seconds when the transaction is allowed to be rolled back.
+    """
+
+    def __init__(self, groupid: int, nodes: Optional[Collection[PgDistNode]], event: str,
+                 timeout: Optional[float] = None, cooldown: Optional[float] = None) -> None:
+        """Create a :class:`PgDistTask` object based on given arguments.
+
+        :param groupid: the groupid from "pg_dist_node".
+        :param nodes: a collection of :class:`PgDistNode` objects that belog to a *groupid*.
+        :param event: an "event" that resulted in creating this task.
+        :param timeout: a transaction timeout if the task resulted in starting a transaction.
+        :param cooldown: the cooldown value for ``citus_update_node()`` UDF call.
+        """
+        super(PgDistTask, self).__init__(groupid, nodes)
+
         # Event that is trying to change or changed the given row.
         # Possible values: before_demote, before_promote, after_promote.
         self.event = event
-        self.nodeid = nodeid
 
         # If transaction was started, we need to COMMIT/ROLLBACK before the deadline
         self.timeout = timeout
@@ -46,25 +339,20 @@ class PgDistNode(object):
         self._event = Event()
 
     def wait(self) -> None:
+        """Wait until this task is processed by a dedicated thread."""
         self._event.wait()
 
     def wakeup(self) -> None:
+        """Notify a thread that created a task that it was processed."""
         self._event.set()
 
     def __eq__(self, other: Any) -> bool:
-        return isinstance(other, PgDistNode) and self.event == other.event\
-            and self.host == other.host and self.port == other.port
+        return isinstance(other, PgDistTask) and self.event == other.event\
+            and super(PgDistTask, self).equals(other)
 
     def __ne__(self, other: Any) -> bool:
         return not self == other
 
-    def __str__(self) -> str:
-        return ('PgDistNode(nodeid={0},group={1},host={2},port={3},event={4})'
-                .format(self.nodeid, self.group, self.host, self.port, self.event))
-
-    def __repr__(self) -> str:
-        return str(self)
-
 
 class Citus(AbstractMPP):
 
@@ -109,11 +397,11 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
             self._connection = postgresql.connection_pool.get(
                 'citus', {'dbname': config['database'],
                           'options': '-c statement_timeout=0 -c idle_in_transaction_session_timeout=0'})
-        self._pg_dist_node: Dict[int, PgDistNode] = {}  # Cache of pg_dist_node: {groupid: PgDistNode()}
-        self._tasks: List[PgDistNode] = []  # Requests to change pg_dist_node, every task is a `PgDistNode`
-        self._in_flight: Optional[PgDistNode] = None  # Reference to the `PgDistNode` being changed in a transaction
-        self._schedule_load_pg_dist_node = True  # Flag that "pg_dist_node" should be queried from the database
-        self._condition = Condition()  # protects _pg_dist_node, _tasks, _in_flight, and _schedule_load_pg_dist_node
+        self._pg_dist_group: Dict[int, PgDistTask] = {}  # Cache of pg_dist_node: {groupid: PgDistTask()}
+        self._tasks: List[PgDistTask] = []  # Requests to change pg_dist_group, every task is a `PgDistTask`
+        self._in_flight: Optional[PgDistTask] = None  # Reference to the `PgDistTask` being changed in a transaction
+        self._schedule_load_pg_dist_group = True  # Flag that "pg_dist_group" should be queried from the database
+        self._condition = Condition()  # protects _pg_dist_group, _tasks, _in_flight, and _schedule_load_pg_dist_group
         self.schedule_cache_rebuild()
 
     def schedule_cache_rebuild(self) -> None:
@@ -122,12 +410,12 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
         Is called to notify handler that it has to refresh its metadata cache from the database.
         """
         with self._condition:
-            self._schedule_load_pg_dist_node = True
+            self._schedule_load_pg_dist_group = True
 
     def on_demote(self) -> None:
         with self._condition:
-            self._pg_dist_node.clear()
-            empty_tasks: List[PgDistNode] = []
+            self._pg_dist_group.clear()
+            empty_tasks: List[PgDistTask] = []
             self._tasks[:] = empty_tasks
             self._in_flight = None
 
@@ -143,22 +431,28 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
             self.schedule_cache_rebuild()
             raise e
 
-    def load_pg_dist_node(self) -> bool:
+    def load_pg_dist_group(self) -> bool:
         """Read from the `pg_dist_node` table and put it into the local cache"""
 
         with self._condition:
-            if not self._schedule_load_pg_dist_node:
+            if not self._schedule_load_pg_dist_group:
                 return True
-            self._schedule_load_pg_dist_node = False
+            self._schedule_load_pg_dist_group = False
 
         try:
-            rows = self.query("SELECT nodeid, groupid, nodename, nodeport, noderole"
-                              " FROM pg_catalog.pg_dist_node WHERE noderole = 'primary'")
+            rows = self.query('SELECT groupid, nodename, nodeport, noderole, nodeid FROM pg_catalog.pg_dist_node')
         except Exception:
             return False
 
+        pg_dist_group: Dict[int, PgDistTask] = {}
+
+        for row in rows:
+            if row[0] not in pg_dist_group:
+                pg_dist_group[row[0]] = PgDistTask(row[0], nodes=set(), event='after_promote')
+            pg_dist_group[row[0]].add(PgDistNode(*row[1:]))
+
         with self._condition:
-            self._pg_dist_node = {r[1]: PgDistNode(r[1], r[2], r[3], 'after_promote', r[0]) for r in rows}
+            self._pg_dist_group = pg_dist_group
         return True
 
     def sync_meta_data(self, cluster: Cluster) -> None:
@@ -166,7 +460,7 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
 
         We can't always rely on REST API calls from worker nodes in order
         to maintain `pg_dist_node`, therefore at least once per heartbeat
-        loop we make sure that workes registered in `self._pg_dist_node`
+        loop we make sure that workes registered in `self._pg_dist_group`
         cache are matching the cluster view from DCS by creating tasks
         the same way as it is done from the REST API."""
 
@@ -177,20 +471,21 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
             if not self.is_alive():
                 self.start()
 
-        self.add_task('after_promote', CITUS_COORDINATOR_GROUP_ID, self._postgresql.connection_string)
+        self.add_task('after_promote', CITUS_COORDINATOR_GROUP_ID, cluster,
+                      self._postgresql.name, self._postgresql.connection_string)
 
-        for group, worker in cluster.workers.items():
+        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('state') == 'running':
-                self.add_task('after_promote', group, leader.conn_url)
+                self.add_task('after_promote', groupid, worker, leader.name, leader.conn_url)
 
-    def find_task_by_group(self, group: int) -> Optional[int]:
+    def find_task_by_groupid(self, groupid: int) -> Optional[int]:
         for i, task in enumerate(self._tasks):
-            if task.group == group:
+            if task.groupid == groupid:
                 return i
 
-    def pick_task(self) -> Tuple[Optional[int], Optional[PgDistNode]]:
+    def pick_task(self) -> Tuple[Optional[int], Optional[PgDistTask]]:
         """Returns the tuple(i, task), where `i` - is the task index in the self._tasks list
 
         Tasks are picked by following priorities:
@@ -198,44 +493,56 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
         1. If there is already a transaction in progress, pick a task
            that that will change already affected worker primary.
         2. If the coordinator address should be changed - pick a task
-           with group=0 (coordinators are always in group 0).
+           with groupid=0 (coordinators are always in groupid 0).
         3. Pick a task that is the oldest (first from the self._tasks)
         """
 
         with self._condition:
             if self._in_flight:
-                i = self.find_task_by_group(self._in_flight.group)
+                i = self.find_task_by_groupid(self._in_flight.groupid)
             else:
                 while True:
-                    i = self.find_task_by_group(CITUS_COORDINATOR_GROUP_ID)  # set_coordinator
+                    i = self.find_task_by_groupid(CITUS_COORDINATOR_GROUP_ID)  # set_coordinator
                     if i is None and self._tasks:
                         i = 0
                     if i is None:
                         break
                     task = self._tasks[i]
-                    if task == self._pg_dist_node.get(task.group):
-                        self._tasks.pop(i)  # nothing to do because cached version of pg_dist_node already matches
+                    if task == self._pg_dist_group.get(task.groupid):
+                        self._tasks.pop(i)  # nothing to do because cached version of pg_dist_group already matches
                     else:
                         break
             task = self._tasks[i] if i is not None else None
 
-            # When tasks are added it could happen that self._pg_dist_node
-            # wasn't ready (self._schedule_load_pg_dist_node is False)
-            # and hence the nodeid wasn't filled.
-            if task and task.group in self._pg_dist_node:
-                task.nodeid = self._pg_dist_node[task.group].nodeid
             return i, task
 
-    def update_node(self, task: PgDistNode) -> None:
-        if task.nodeid is not None:
+    def update_node(self, groupid: int, node: PgDistNode, cooldown: float = 10000) -> None:
+        if node.role not in ('primary', 'secondary', 'demoted'):
+            self.query('SELECT pg_catalog.citus_remove_node(%s, %s)', node.host, node.port)
+        elif node.nodeid is not None:
+            host = node.host + ('-demoted' if node.role == 'demoted' else '')
             self.query('SELECT pg_catalog.citus_update_node(%s, %s, %s, true, %s)',
-                       task.nodeid, task.host, task.port, task.cooldown)
-        elif task.event != 'before_demote':
-            task.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
-                                     task.host, task.port, task.group)[0][0]
+                       node.nodeid, host, node.port, cooldown)
+        elif node.role != 'demoted':
+            node.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, %s, 'default')",
+                                     node.host, node.port, groupid, node.role)[0][0]
 
-    def process_task(self, task: PgDistNode) -> bool:
-        """Updates a single row in `pg_dist_node` table, optionally in a transaction.
+    def update_group(self, task: PgDistTask, transaction: bool) -> None:
+        current_state = self._in_flight\
+            or self._pg_dist_group.get(task.groupid)\
+            or PgDistTask(task.groupid, set(), 'after_promote')
+        transitions = list(task.transition(current_state))
+        if transitions:
+            if not transaction and len(transitions) > 1:
+                self.query('BEGIN')
+            for node in transitions:
+                self.update_node(task.groupid, node, task.cooldown)
+            if not transaction and len(transitions) > 1:
+                task.failover = False
+                self.query('COMMIT')
+
+    def process_task(self, task: PgDistTask) -> bool:
+        """Updates a single row in `pg_dist_group` table, optionally in a transaction.
 
         The transaction is started if we do a demote of the worker node or before promoting the other worker if
         there is no transaction in progress. And, the transaction is committed when the switchover/failover completed.
@@ -246,34 +553,30 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
         .. note:
             Read access to `self._in_flight` isn't protected because we know it can't be changed outside of our thread.
 
-        :param task: reference to a :class:`PgDistNode` object that represents a row to be updated/created.
-        :returns: `True` if the row was succesfully created/updated or transaction in progress
-            was committed as an indicator that the `self._pg_dist_node` cache should be updated,
+        :param task: reference to a :class:`PgDistTask` object that represents a row to be updated/created.
+        :returns: ``True`` if the row was succesfully created/updated or transaction in progress
+            was committed as an indicator that the `self._pg_dist_group` cache should be updated,
             or, if the new transaction was opened, this method returns `False`.
         """
 
         if task.event == 'after_promote':
-            # The after_promote may happen without previous before_demote and/or
-            # before_promore.  In this case we just call self.update_node() method.
-            # If there is a transaction in progress, it could be that it already did
-            # required changes and we can simply COMMIT.
-            if not self._in_flight or self._in_flight.host != task.host or self._in_flight.port != task.port:
-                self.update_node(task)
+            self.update_group(task, self._in_flight is not None)
             if self._in_flight:
                 self.query('COMMIT')
+            task.failover = False
             return True
         else:  # before_demote, before_promote
             if task.timeout:
                 task.deadline = time.time() + task.timeout
             if not self._in_flight:
                 self.query('BEGIN')
-            self.update_node(task)
+            self.update_group(task, True)
         return False
 
     def process_tasks(self) -> None:
         while True:
             # Read access to `_in_flight` isn't protected because we know it can't be changed outside of our thread.
-            if not self._in_flight and not self.load_pg_dist_node():
+            if not self._in_flight and not self.load_pg_dist_group():
                 break
 
             i, task = self.pick_task()
@@ -287,7 +590,7 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
             with self._condition:
                 if self._tasks:
                     if update_cache:
-                        self._pg_dist_node[task.group] = task
+                        self._pg_dist_group[task.groupid] = task
 
                     if update_cache is False:  # an indicator that process_tasks has started a transaction
                         self._in_flight = task
@@ -302,7 +605,7 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
         while True:
             try:
                 with self._condition:
-                    if self._schedule_load_pg_dist_node:
+                    if self._schedule_load_pg_dist_group:
                         timeout = -1
                     elif self._in_flight:
                         timeout = self._in_flight.deadline - time.time() if self._tasks else None
@@ -319,9 +622,9 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
             except Exception:
                 logger.exception('run')
 
-    def _add_task(self, task: PgDistNode) -> bool:
+    def _add_task(self, task: PgDistTask) -> bool:
         with self._condition:
-            i = self.find_task_by_group(task.group)
+            i = self.find_task_by_groupid(task.groupid)
 
             # The `PgDistNode.timeout` == None is an indicator that it was scheduled from the sync_meta_data().
             if task.timeout is None:
@@ -333,37 +636,48 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
                 # key is updated in DCS. Therefore it is possible that :func:`sync_meta_data` will try to create a task
                 # based on the outdated values of "state"/"role". To solve it we introduce an artificial timeout.
                 # Only when the timeout is reached new tasks could be scheduled from sync_meta_data()
-                if self._in_flight and self._in_flight.group == task.group and self._in_flight.timeout is not None\
+                if self._in_flight and self._in_flight.groupid == task.groupid and self._in_flight.timeout is not None\
                         and self._in_flight.deadline > time.time():
                     return False
 
-            # Override already existing task for the same worker group
+            # Override already existing task for the same worker groupid
             if i is not None:
                 if task != self._tasks[i]:
                     logger.debug('Overriding existing task: %s != %s', self._tasks[i], task)
                     self._tasks[i] = task
                     self._condition.notify()
                     return True
-            # Add the task to the list if Worker node state is different from the cached `pg_dist_node`
-            elif self._schedule_load_pg_dist_node or task != self._pg_dist_node.get(task.group)\
-                    or self._in_flight and task.group == self._in_flight.group:
+            # Add the task to the list if Worker node state is different from the cached `pg_dist_group`
+            elif self._schedule_load_pg_dist_group or task != self._pg_dist_group.get(task.groupid)\
+                    or self._in_flight and task.groupid == self._in_flight.groupid:
                 logger.debug('Adding the new task: %s', task)
                 self._tasks.append(task)
                 self._condition.notify()
                 return True
         return False
 
-    def add_task(self, event: str, group: int, conn_url: str,
-                 timeout: Optional[float] = None, cooldown: Optional[float] = None) -> Optional[PgDistNode]:
+    @staticmethod
+    def _pg_dist_node(role: str, conn_url: str) -> Optional[PgDistNode]:
         try:
             r = urlparse(conn_url)
+            if r.hostname:
+                return PgDistNode(r.hostname, r.port or 5432, role)
         except Exception as e:
-            return logger.error('Failed to parse connection url %s: %r', conn_url, e)
-        host = r.hostname
-        if host:
-            port = r.port or 5432
-            task = PgDistNode(group, host, port, event, timeout=timeout, cooldown=cooldown)
-            return task if self._add_task(task) else None
+            logger.error('Failed to parse connection url %s: %r', conn_url, e)
+
+    def add_task(self, event: str, groupid: int, cluster: Cluster, leader_name: str, leader_url: str,
+                 timeout: Optional[float] = None, cooldown: Optional[float] = None) -> Optional[PgDistTask]:
+        primary = self._pg_dist_node('demoted' if event == 'before_demote' else 'primary', leader_url)
+        if not primary:
+            return
+
+        task = PgDistTask(groupid, {primary}, event=event, timeout=timeout, cooldown=cooldown)
+        for member in cluster.members:
+            secondary = self._pg_dist_node('secondary', member.conn_url)\
+                if member.name != leader_name and member.is_running and member.conn_url else None
+            if secondary:
+                task.add(secondary)
+        return task if self._add_task(task) else None
 
     def handle_event(self, cluster: Cluster, event: Dict[str, Any]) -> None:
         if not self.is_alive():
@@ -371,10 +685,10 @@ class CitusHandler(Citus, AbstractMPPHandler, Thread):
 
         worker = cluster.workers.get(event['group'])
         if not (worker and worker.leader and worker.leader.name == event['leader'] and worker.leader.conn_url):
-            return
+            return logger.info('Discarding event %s', event)
 
-        task = self.add_task(event['type'], event['group'],
-                             worker.leader.conn_url,
+        task = self.add_task(event['type'], event['group'], worker,
+                             worker.leader.name, worker.leader.conn_url,
                              event['timeout'], event['cooldown'] * 1000)
         if task and event['type'] == 'before_demote':
             task.wait()

patroni/version.py

@@ -2,4 +2,4 @@
 
 :var __version__: the current Patroni version.
 """
-__version__ = '3.3.0'
+__version__ = '3.2.2'

tests/__init__.py

@@ -179,8 +179,12 @@ class MockCursor(object):
                                  b'3\t0/403DD98\tno recovery target specified\n')]
         elif sql.startswith('SELECT pg_catalog.citus_add_node'):
             self.results = [(2,)]
-        elif sql.startswith('SELECT nodeid, groupid'):
-            self.results = [(1, 0, 'host1', 5432, 'primary'), (2, 1, 'host2', 5432, 'primary')]
+        elif sql.startswith('SELECT groupid, nodename'):
+            self.results = [(0, 'host1', 5432, 'primary', 1),
+                            (0, '127.0.0.1', 5436, 'secondary', 2),
+                            (1, 'host4', 5432, 'primary', 3),
+                            (1, '127.0.0.1', 5437, 'secondary', 4),
+                            (1, '127.0.0.1', 5438, 'secondary', 5)]
         else:
             self.results = [(None, None, None, None, None, None, None, None, None, None)]
         self.rowcount = len(self.results)

tests/test_citus.py

@@ -1,6 +1,10 @@
 import time
+import unittest
+
+from copy import deepcopy
 from mock import Mock, patch, PropertyMock
-from patroni.postgresql.mpp.citus import CitusHandler
+from typing import List
+from patroni.postgresql.mpp.citus import CitusHandler, PgDistGroup, PgDistNode
 from patroni.psycopg import ProgrammingError
 
 from . import BaseTestPostgresql, MockCursor, psycopg_connect, SleepException
@@ -20,7 +24,7 @@ class TestCitus(BaseTestPostgresql):
     @patch('time.time', Mock(side_effect=[100, 130, 160, 190, 220, 250, 280, 310, 340, 370]))
     @patch('patroni.postgresql.mpp.citus.logger.exception', Mock(side_effect=SleepException))
     @patch('patroni.postgresql.mpp.citus.logger.warning')
-    @patch('patroni.postgresql.mpp.citus.PgDistNode.wait', Mock())
+    @patch('patroni.postgresql.mpp.citus.PgDistTask.wait', Mock())
     @patch.object(CitusHandler, 'is_alive', Mock(return_value=True))
     def test_run(self, mock_logger_warning):
         # `before_demote` or `before_promote` REST API calls starting a
@@ -32,11 +36,11 @@ class TestCitus(BaseTestPostgresql):
 
         self.c.handle_event(self.cluster, {'type': 'before_demote', 'group': 1,
                                            'leader': 'leader', 'timeout': 30, 'cooldown': 10})
-        self.c.add_task('after_promote', 2, 'postgres://host3:5432/postgres')
+        self.c.add_task('after_promote', 2, self.cluster, self.cluster.leader_name, 'postgres://host3:5432/postgres')
         self.assertRaises(SleepException, self.c.run)
         mock_logger_warning.assert_called_once()
         self.assertTrue(mock_logger_warning.call_args[0][0].startswith('Rolling back transaction'))
-        self.assertTrue(repr(mock_logger_warning.call_args[0][1]).startswith('PgDistNode'))
+        self.assertTrue(repr(mock_logger_warning.call_args[0][1]).startswith('PgDistTask'))
 
     @patch.object(CitusHandler, 'is_alive', Mock(return_value=False))
     @patch.object(CitusHandler, 'start', Mock())
@@ -54,59 +58,68 @@ class TestCitus(BaseTestPostgresql):
     def test_add_task(self):
         with patch('patroni.postgresql.mpp.citus.logger.error') as mock_logger, \
                 patch('patroni.postgresql.mpp.citus.urlparse', Mock(side_effect=Exception)):
-            self.c.add_task('', 1, None)
+            self.c.add_task('', 1, self.cluster, '', None)
             mock_logger.assert_called_once()
 
         with patch('patroni.postgresql.mpp.citus.logger.debug') as mock_logger:
-            self.c.add_task('before_demote', 1, 'postgres://host:5432/postgres', 30)
+            self.c.add_task('before_demote', 1, self.cluster,
+                            self.cluster.leader_name, 'postgres://host:5432/postgres', 30)
             mock_logger.assert_called_once()
             self.assertTrue(mock_logger.call_args[0][0].startswith('Adding the new task:'))
 
         with patch('patroni.postgresql.mpp.citus.logger.debug') as mock_logger:
-            self.c.add_task('before_promote', 1, 'postgres://host:5432/postgres', 30)
+            self.c.add_task('before_promote', 1, self.cluster,
+                            self.cluster.leader_name, 'postgres://host:5432/postgres', 30)
             mock_logger.assert_called_once()
             self.assertTrue(mock_logger.call_args[0][0].startswith('Overriding existing task:'))
 
-        # add_task called from sync_meta_data should not override already scheduled or in flight task until deadline
-        self.assertIsNotNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres', 30))
-        self.assertIsNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
+        # add_task called from sync_pg_dist_node should not override already scheduled or in flight task until deadline
+        self.assertIsNotNone(self.c.add_task('after_promote', 1, self.cluster,
+                                             self.cluster.leader_name, 'postgres://host:5432/postgres', 30))
+        self.assertIsNone(self.c.add_task('after_promote', 1, self.cluster,
+                                          self.cluster.leader_name, 'postgres://host:5432/postgres'))
         self.c._in_flight = self.c._tasks.pop()
         self.c._in_flight.deadline = self.c._in_flight.timeout + time.time()
-        self.assertIsNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
+        self.assertIsNone(self.c.add_task('after_promote', 1, self.cluster,
+                                          self.cluster.leader_name, 'postgres://host:5432/postgres'))
         self.c._in_flight.deadline = 0
-        self.assertIsNotNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
+        self.assertIsNotNone(self.c.add_task('after_promote', 1, self.cluster,
+                                             self.cluster.leader_name, 'postgres://host:5432/postgres'))
 
         # If there is no transaction in progress and cached pg_dist_node matching desired state task should not be added
         self.c._schedule_load_pg_dist_node = False
-        self.c._pg_dist_node[self.c._in_flight.group] = self.c._in_flight
+        self.c._pg_dist_group[self.c._in_flight.groupid] = self.c._in_flight
         self.c._in_flight = None
-        self.assertIsNone(self.c.add_task('after_promote', 1, 'postgres://host:5432/postgres'))
+        self.assertIsNone(self.c.add_task('after_promote', 1, self.cluster,
+                                          self.cluster.leader_name, 'postgres://host:5432/postgres'))
 
     def test_pick_task(self):
-        self.c.add_task('after_promote', 1, 'postgres://host2:5432/postgres')
-        with patch.object(CitusHandler, 'process_task') as mock_process_task:
+        self.c.add_task('after_promote', 0, self.cluster, self.cluster.leader_name, 'postgres://host1:5432/postgres')
+        with patch.object(CitusHandler, 'update_node') as mock_update_node:
             self.c.process_tasks()
-            # process_task() shouln't be called because pick_task double checks with _pg_dist_node
-            mock_process_task.assert_not_called()
+            # process_task() shouln't be called because pick_task double checks with _pg_dist_group
+            mock_update_node.assert_not_called()
 
     def test_process_task(self):
-        self.c.add_task('after_promote', 0, 'postgres://host2:5432/postgres')
-        task = self.c.add_task('before_promote', 1, 'postgres://host4:5432/postgres', 30)
+        self.c.add_task('after_promote', 1, self.cluster, self.cluster.leader_name, 'postgres://host2:5432/postgres')
+        task = self.c.add_task('before_promote', 1, self.cluster,
+                               self.cluster.leader_name, 'postgres://host4:5432/postgres', 30)
         self.c.process_tasks()
         self.assertTrue(task._event.is_set())
 
         # the after_promote should result only in COMMIT
-        task = self.c.add_task('after_promote', 1, 'postgres://host4:5432/postgres', 30)
+        task = self.c.add_task('after_promote', 1, self.cluster,
+                               self.cluster.leader_name, 'postgres://host4:5432/postgres', 30)
         with patch.object(CitusHandler, 'query') as mock_query:
             self.c.process_tasks()
             mock_query.assert_called_once()
             self.assertEqual(mock_query.call_args[0][0], 'COMMIT')
 
     def test_process_tasks(self):
-        self.c.add_task('after_promote', 0, 'postgres://host2:5432/postgres')
+        self.c.add_task('after_promote', 0, self.cluster, self.cluster.leader_name, 'postgres://host2:5432/postgres')
         self.c.process_tasks()
 
-        self.c.add_task('after_promote', 0, 'postgres://host3:5432/postgres')
+        self.c.add_task('after_promote', 0, self.cluster, self.cluster.leader_name, 'postgres://host3:5432/postgres')
         with patch('patroni.postgresql.mpp.citus.logger.error') as mock_logger, \
                 patch.object(CitusHandler, 'query', Mock(side_effect=Exception)):
             self.c.process_tasks()
@@ -118,16 +131,17 @@ class TestCitus(BaseTestPostgresql):
 
     @patch('patroni.postgresql.mpp.citus.logger.error')
     @patch.object(MockCursor, 'execute', Mock(side_effect=Exception))
-    def test_load_pg_dist_node(self, mock_logger):
-        # load_pg_dist_node() triggers, query fails and exception is property handled
+    def test_load_pg_dist_group(self, mock_logger):
+        # load_pg_dist_group) triggers, query fails and exception is property handled
         self.c.process_tasks()
-        self.assertTrue(self.c._schedule_load_pg_dist_node)
+        self.assertTrue(self.c._schedule_load_pg_dist_group)
         mock_logger.assert_called_once()
         self.assertTrue(mock_logger.call_args[0][0].startswith('Exception when executing query'))
-        self.assertTrue(mock_logger.call_args[0][1].startswith('SELECT nodeid, groupid, '))
+        self.assertTrue(mock_logger.call_args[0][1].startswith('SELECT groupid, nodename, '))
 
     def test_wait(self):
-        task = self.c.add_task('before_demote', 1, 'postgres://host:5432/postgres', 30)
+        task = self.c.add_task('before_demote', 1, self.cluster,
+                               self.cluster.leader_name, u'postgres://host:5432/postgres', 30)
         task._event.wait = Mock()
         task.wait()
 
@@ -171,3 +185,209 @@ class TestCitus(BaseTestPostgresql):
             self.c.bootstrap()
         mock_logger.assert_called_once()
         self.assertTrue(mock_logger.call_args[0][0].startswith('Exception when creating database'))
+
+
+class TestGroupTransition(unittest.TestCase):
+    nodeid = 100
+
+    def map_to_sql(self, group: int, transition: PgDistNode) -> str:
+        if transition.role not in ('primary', 'demoted', 'secondary'):
+            return "citus_remove_node('{0}', {1})".format(transition.host, transition.port)
+        elif transition.nodeid:
+            host = transition.host + ('-demoted' if transition.role == 'demoted' else '')
+            return "citus_update_node({0}, '{1}', {2})".format(transition.nodeid, host, transition.port)
+        else:
+            transition.nodeid = self.nodeid
+            self.nodeid += 1
+
+            return "citus_add_node('{0}', {1}, {2}, '{3}')".format(transition.host, transition.port,
+                                                                   group, transition.role)
+
+    def check_transitions(self, old_topology: PgDistGroup, new_topology: PgDistGroup,
+                          expected_transitions: List[str]) -> None:
+        check_topology = deepcopy(old_topology)
+
+        transitions: List[str] = []
+        for node in new_topology.transition(old_topology):
+            self.assertTrue(node not in check_topology or (check_topology.get(node) or node).role == 'demoted')
+            old_node = node.nodeid and next(iter(v for v in check_topology if v.nodeid == node.nodeid), None)
+            if old_node:
+                check_topology.discard(old_node)
+            transitions.append(self.map_to_sql(new_topology.groupid, node))
+            check_topology.add(node)
+        self.assertEqual(transitions, expected_transitions)
+
+    def test_new_topology(self):
+        old = PgDistGroup(0)
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'primary'),
+                              PgDistNode('2', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=100),
+                                   PgDistNode('2', 5432, 'secondary', nodeid=101)})
+        self.check_transitions(old, new,
+                               ["citus_add_node('1', 5432, 0, 'primary')",
+                                "citus_add_node('2', 5432, 0, 'secondary')"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary'),
+                              PgDistNode('2', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('2', 5432, 'primary', nodeid=1),
+                                   PgDistNode('1', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new,
+                               ["citus_update_node(1, '1-demoted', 5432)",
+                                "citus_update_node(2, '1', 5432)",
+                                "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_failover(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('2', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('2', 5432, 'primary', nodeid=1),
+                                   PgDistNode('1', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new,
+                               ["citus_update_node(1, '1-demoted', 5432)",
+                                "citus_update_node(2, '1', 5432)",
+                                "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_failover_and_new_secondary(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('2', 5432, 'primary'),
+                              PgDistNode('3', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('2', 5432, 'primary', nodeid=1),
+                                   PgDistNode('3', 5432, 'secondary', nodeid=2)})
+        # the secondary record is used to add the new standby and primary record is updated with the new hostname
+        self.check_transitions(old, new, ["citus_update_node(2, '3', 5432)", "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_and_new_secondary_primary_gone(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('2', 5432, 'primary'),
+                              PgDistNode('3', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('2', 5432, 'primary', nodeid=1),
+                                   PgDistNode('3', 5432, 'secondary', nodeid=2)})
+        # the secondary record is used to add the new standby and primary record is updated with the new hostname
+        self.check_transitions(old, new, ["citus_update_node(2, '3', 5432)", "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_secondary_replaced(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'primary'),
+                              PgDistNode('3', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                                   PgDistNode('3', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new, ["citus_update_node(2, '3', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_secondary_repmoved(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2),
+                              PgDistNode('3', 5432, 'secondary', nodeid=3)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'primary'),
+                              PgDistNode('3', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                                   PgDistNode('3', 5432, 'secondary', nodeid=3)})
+        self.check_transitions(old, new, ["citus_remove_node('2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_and_secondary_removed(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2),
+                              PgDistNode('3', 5432, 'secondary', nodeid=3)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary'),
+                              PgDistNode('2', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary', nodeid=2),
+                                   PgDistNode('2', 5432, 'primary', nodeid=1),
+                                   PgDistNode('3', 5432, 'secondary', nodeid=3)})
+        self.check_transitions(old, new,
+                               ["citus_update_node(1, '1-demoted', 5432)",
+                                "citus_update_node(2, '1', 5432)",
+                                "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_and_new_secondary(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary'),
+                              PgDistNode('2', 5432, 'primary'),
+                              PgDistNode('3', 5432, 'secondary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary', nodeid=2),
+                                   PgDistNode('2', 5432, 'primary', nodeid=1)})
+        self.check_transitions(old, new,
+                               ["citus_update_node(1, '1-demoted', 5432)",
+                                "citus_update_node(2, '1', 5432)",
+                                "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_failover_to_new_node_secondary_remains(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('2', 5432, 'secondary'),
+                              PgDistNode('3', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('3', 5432, 'primary', nodeid=1),
+                                   PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new, ["citus_update_node(1, '3', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_failover_to_new_node_secondary_removed(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('3', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('3', 5432, 'primary', nodeid=1),
+                                   PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        # the secondary record needs to be removed before we update the primary record
+        self.check_transitions(old, new, ["citus_update_node(1, '3', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_to_new_node_and_secondary_removed(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary'),
+                              PgDistNode('3', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('3', 5432, 'primary', nodeid=1),
+                                   PgDistNode('1', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new, ["citus_update_node(1, '3', 5432)", "citus_update_node(2, '1', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_with_pause(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'primary', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted', nodeid=1),
+                                   PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        self.check_transitions(old, new, ["citus_update_node(1, '1-demoted', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_after_paused_connections(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('2', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary', nodeid=2),
+                                   PgDistNode('2', 5432, 'primary', nodeid=1)})
+        self.check_transitions(old, new, ["citus_update_node(2, '1', 5432)", "citus_update_node(1, '2', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_to_new_node_after_paused_connections(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('3', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('1', 5432, 'secondary', nodeid=2),
+                                   PgDistNode('3', 5432, 'primary', nodeid=1)})
+        self.check_transitions(old, new, ["citus_update_node(1, '3', 5432)", "citus_update_node(2, '1', 5432)"])
+        self.assertTrue(new.equals(expected, True))
+
+    def test_switchover_to_new_node_after_paused_connections_secondary_added(self):
+        old = PgDistGroup(0, {PgDistNode('1', 5432, 'demoted', nodeid=1),
+                              PgDistNode('2', 5432, 'secondary', nodeid=2)})
+        new = PgDistGroup(0, {PgDistNode('4', 5432, 'secondary'),
+                              PgDistNode('3', 5432, 'primary')})
+        expected = PgDistGroup(0, {PgDistNode('4', 5432, 'secondary', nodeid=2),
+                                   PgDistNode('3', 5432, 'primary', nodeid=1)})
+        self.check_transitions(old, new, ["citus_update_node(1, '3', 5432)", "citus_update_node(2, '4', 5432)"])
+        self.assertTrue(new.equals(expected, True))

tests/test_consul.py

@@ -160,7 +160,7 @@ class TestConsul(unittest.TestCase):
         self.c.set_ttl(20)
         self.c._do_refresh_session = Mock()
         self.assertFalse(self.c.take_leader())
-        with patch('time.time', Mock(side_effect=[0, 0, 0, 100, 100, 100])):
+        with patch('time.time', Mock(side_effect=[0, 0, 0, 100, 100])):
             self.assertFalse(self.c.take_leader())
 
     @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))

tests/test_etcd3.py

@@ -273,7 +273,7 @@ class TestEtcd3(BaseTestEtcd3):
 
     def test_attempt_to_acquire_leader(self):
         self.assertFalse(self.etcd3.attempt_to_acquire_leader())
-        with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 100, 200, 300])):
+        with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 100, 200])):
             self.assertFalse(self.etcd3.attempt_to_acquire_leader())
         with patch('time.time', Mock(side_effect=[0, 100, 200, 300, 400])):
             self.assertRaises(Etcd3Error, self.etcd3.attempt_to_acquire_leader)