Register Citus secondaries in pg_dist_node

1. All nodes with role == 'replica' and state == 'running' are
   are registered. In case is state isn't running the node is removed.
2. In case of failover/switchover we always first update the primary
3. When switching to a registered secondary we call citus_update_node()
   three times: rename primary to primary-demoted, put the primary name
   to a promoted secondary row and put the promoted secondary name to
   the primary row

State transitions are produced by the transition() method. First of all
the method makes sure that the actual primary is registered in the
metadata. In case if for a given group the primary didn't change, the
method registers new secondaries and removes secondaries that are gone.
It prefers to use citus_update_node() UDF to replace gone secondaries
with added.

Communication protocol between primary nodes remains the same and all
old features work without any changes.

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/postgresql/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 .connection import Connection
 from ..dcs import CITUS_COORDINATOR_GROUP_ID, Cluster
@@ -19,19 +19,296 @@ 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.
+
+        We use (*host*, *port*) tuple here because it is one of the UNIQUE constraints on the "pg_dist_node" table.
+        The *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.
+
+        It is used exclusively to support overriden :func:``PgDistNode.__hash__``.
+
+        :returns: ``True`` if *host* and *port* between two instances are the same.
+        """
+        return isinstance(other, PgDistNode) and self.host == other.host and self.port == other.port
+
+    def __ne__(self, other: Any) -> bool:
+        return not self == other
+
+    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 "primary".
+        """
+        return self.role in ('primary', 'demoted')
+
+
+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
+    - failover and switchover
+
+    Typically there will be at least one :class:`PgDistNode` object registered ('primary').
+    In adding to that there could be some "secondaries".
+
+    :ivar failover: whether as a result of :func:`transition` method call the "primary" row should be updated.
+    :ivar group: the "groupid" from "pg_dist_node"
+    """
+
+    def __init__(self, group: int, nodes: Optional[Collection[PgDistNode]] = None) -> None:
+        """Creates a :class:`PgDistGroup` object based on given arguments.
+
+        :param group: the groupid from "pg_dist_node".
+        :param nodes: a collection of :class:`PgDistNode` objects that belog to a *group*.
+        """
+        self.failover = False
+        self.group = group
+
+        if nodes:
+            self.update(nodes)
+
+    @staticmethod
+    def _node_hash(node: PgDistNode, include_nodeid: bool = False) -> Tuple[str, int, str, Optional[int]]:
+        """Helper function to compare two :class:`PgDistGroup` objects.
+
+        .. note::
+
+            *include_nodeid* is set to `True` only in unit-tests.
+
+        :param node: the PgDistNode we want to build hash for.
+        :param include_nodeid: whether *nodeid* should be taken into account when comparison is performed.
+        :returns: :class:`tuple` object with *host*, *port*, *role*, and optionally *nodeid*
+        """
+        return node.host, node.port, node.role, (node.nodeid if include_nodeid else None)
+
+    def equals(self, other: 'PgDistGroup', check_nodeid: bool = False) -> bool:
+        """Compares two :class:`PgDistGroup` objects.
+
+        .. note::
+
+            Normally only *host*, *port*, and *role* values are compared for all :class:`PgDistNode` objects.
+            But, optionally it can also compare *nodeid* if *check_nodeid* if set to `True` (used only in unit-tests).
+
+        :param other: what we want to compare with.
+        :param check_nodeid: whether *nodeid* should be compared in addition to *host*, *port*, and *role*.
+        :returns: `True` if all two objects are identical.
+        """
+        return set(self._node_hash(v, check_nodeid) for v in self)\
+            == set(self._node_hash(v, check_nodeid) for v in other)
+
+    def primary(self) -> Optional[PgDistNode]:
+        """Finds and returns :class:`PgDistNode` object that represents "primary"."""
+        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 given set.
+
+        .. note::
+            It is necessary because :func:`__hash__` and :func:`__eq__` methods in :class:`PgDistNode`
+            are redefined and effectively they check only *host* and *port* attributes.
+
+        :param value: the key we search for.
+        :returns: the actual value from this :class:`set` object.
+        """
+        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::
+
+            In addition to the yielding transactions this method fills up *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".
+            - updating "broken" nodes always works and metadata is synced asynchnonously after commit.
+
+        Following these rules below is an example of the switchover between node1 (primary) and node2 (secondary).
+
+        :Example:
+
+            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 know topology registered in "pg_dist_node" for a given *group*
+        :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:
+            yield new_primary
+        elif old_primary == new_primary:
+            new_primary.nodeid = old_primary.nodeid
+            # Controlled switchover with pausing client connections.
+            # Achived 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 gone away.
+            # Since nodes can't be removed while metadata isn't synced we have to temporary "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 temporary "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
+
+        # Remove remaining nodes that are gone, but only in case if metadata is in sync
+        for g in gone_nodes:
+            if self.failover:  # Otherwise add these nodes to the new topology
+                self.add(g)
+            else:
+                yield PgDistNode(g.host, g.port, '')
+
+        # Add new nodes to the metadata, but only in case if metadata is in sync
+        for a in added_nodes:
+            if self.failover:
+                self.discard(a)  # Otherwise remove them from the new topology
+            else:
+                yield a
+
+
+class PgDistTask(PgDistGroup):
+    """A "task" that represents the current or desired state of "pg_dist_node" for a provided *group*.
+
+    :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, group: 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 group: the groupid from "pg_dist_node".
+        :param nodes: a collection of :class:`PgDistNode` objects that belog to a *group*.
+        :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__(group, 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 +323,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 CitusHandler(Thread):
 
@@ -74,11 +346,11 @@ class CitusHandler(Thread):
         self._postgresql = postgresql
         self._config = config
         self._connection = Connection()
-        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 is_enabled(self) -> bool:
@@ -101,11 +373,11 @@ class CitusHandler(Thread):
 
     def schedule_cache_rebuild(self) -> None:
         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()
+            self._pg_dist_group.clear()
             self._tasks[:] = []
             self._in_flight = None
 
@@ -123,22 +395,28 @@ class CitusHandler(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:
-            cursor = self.query("SELECT nodeid, groupid, nodename, nodeport, noderole"
-                                " FROM pg_catalog.pg_dist_node WHERE noderole = 'primary'")
+            cursor = 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 cursor:
+            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 cursor}
+            self._pg_dist_group = pg_dist_group
         return True
 
     def sync_pg_dist_node(self, cluster: Cluster) -> None:
@@ -146,7 +424,7 @@ class CitusHandler(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."""
 
@@ -157,20 +435,21 @@ class CitusHandler(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():
             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', group, worker, leader.name, leader.conn_url)
 
     def find_task_by_group(self, group: int) -> Optional[int]:
         for i, task in enumerate(self._tasks):
             if task.group == group:
                 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:
@@ -191,31 +470,43 @@ class CitusHandler(Thread):
                     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.group):
+                        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, group: 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':
-            row = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
-                             task.host, task.port, task.group).fetchone()
+                       node.nodeid, host, node.port, cooldown)
+        elif node.role != 'demoted':
+            row = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, %s, 'default')",
+                             node.host, node.port, group, node.role).fetchone()
             if row is not None:
-                task.nodeid = row[0]
+                node.nodeid = row[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.group)\
+            or PgDistTask(task.group, 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.group, 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.
@@ -226,34 +517,30 @@ class CitusHandler(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.
+        :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_node` cache should be updated,
+            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()
@@ -267,7 +554,7 @@ class CitusHandler(Thread):
             with self._condition:
                 if self._tasks:
                     if update_cache:
-                        self._pg_dist_node[task.group] = task
+                        self._pg_dist_group[task.group] = task
 
                     if update_cache is False:  # an indicator that process_tasks has started a transaction
                         self._in_flight = task
@@ -282,7 +569,7 @@ class CitusHandler(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
@@ -299,20 +586,20 @@ class CitusHandler(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)
 
-            # The `PgDistNode.timeout` == None is an indicator that it was scheduled from the sync_pg_dist_node().
+            # The `PgDistTask.timeout` == None is an indicator that it was scheduled from the sync_pg_dist_group().
             if task.timeout is None:
                 # We don't want to override the already existing task created from REST API.
                 if i is not None and self._tasks[i].timeout is not None:
                     return False
 
                 # There is a little race condition with tasks created from REST API - the call made "before" the member
-                # key is updated in DCS. Therefore it is possible that :func:`sync_pg_dist_node` will try to create a
+                # key is updated in DCS. Therefore it is possible that :func:`sync_pg_dist_group` 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_pg_dist_node()
+                # Only when the timeout is reached new tasks could be scheduled from sync_pg_dist_group()
                 if self._in_flight and self._in_flight.group == task.group and self._in_flight.timeout is not None\
                         and self._in_flight.deadline > time.time():
                     return False
@@ -324,8 +611,8 @@ class CitusHandler(Thread):
                     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)\
+            # 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.group)\
                     or self._in_flight and task.group == self._in_flight.group:
                 logger.debug('Adding the new task: %s', task)
                 self._tasks.append(task)
@@ -333,17 +620,28 @@ class CitusHandler(Thread):
                 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, group: 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(group, {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():
@@ -351,10 +649,10 @@ class CitusHandler(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()

tests/__init__.py

@@ -137,8 +137,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)]
 

tests/test_citus.py

@@ -1,6 +1,10 @@
 import time
+import unittest
+
+from copy import deepcopy
 from mock import Mock, patch
-from patroni.postgresql.citus import CitusHandler
+from typing import List
+from patroni.postgresql.citus import CitusHandler, PgDistGroup, PgDistNode
 
 from . import BaseTestPostgresql, MockCursor, psycopg_connect, SleepException
 from .test_ha import get_cluster_initialized_with_leader
@@ -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.citus.logger.exception', Mock(side_effect=SleepException))
     @patch('patroni.postgresql.citus.logger.warning')
-    @patch('patroni.postgresql.citus.PgDistNode.wait', Mock())
+    @patch('patroni.postgresql.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.citus.logger.error') as mock_logger,\
                 patch('patroni.postgresql.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.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.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_pg_dist_node 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'))
+        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.group] = 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.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.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()
 
@@ -161,3 +175,209 @@ class TestCitus(BaseTestPostgresql):
                                                          'type': 'logical', 'database': 'citus', 'plugin': 'pgoutput'}))
         self.assertTrue(self.c.ignore_replication_slot({'name': 'citus_shard_split_slot_1_2_3',
                                                         'type': 'logical', 'database': 'citus', 'plugin': 'citus'}))
+
+
+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.group, 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))