[MISC] Add back raft encryption check on upgrade

src/cluster.py

@@ -459,7 +459,7 @@ def is_creating_backup(self) -> bool:
             for member in r.json()["members"]
         )
 
-    def is_replication_healthy(self) -> bool:
+    def is_replication_healthy(self, raft_encryption: bool = False) -> bool:
         """Return whether the replication is healthy."""
         try:
             for attempt in Retrying(stop=stop_after_delay(60), wait=wait_fixed(3)):
@@ -477,6 +477,12 @@ def is_replication_healthy(self) -> bool:
                             auth=self._patroni_auth,
                             timeout=PATRONI_TIMEOUT,
                         )
+                        # If raft is getting encrypted some of the calls will fail
+                        if member_status.status_code == 503 and raft_encryption:
+                            logger.warning(
+                                f"Failed replication check for {members_ip} during raft encryption"
+                            )
+                            continue
                         if member_status.status_code != 200:
                             logger.debug(
                                 f"Failed replication check for {members_ip} with code {member_status.status_code}"

src/upgrade.py

@@ -157,6 +157,14 @@ def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
             self.set_unit_failed()
             return
 
+        raft_encryption = (
+            int(
+                json.loads(self.peer_relation.data[self.charm.app].get("dependencies", "{}"))
+                .get("charm", {})
+                .get("version", 0)
+            )
+            < 3
+        )
         self.charm._setup_exporter()
         self.charm.backup.start_stop_pgbackrest_service()
 
@@ -178,7 +186,7 @@ def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
                         not self.charm._patroni.member_started
                         or self.charm.unit.name.replace("/", "-")
                         not in self.charm._patroni.cluster_members
-                        or not self.charm._patroni.is_replication_healthy()
+                        or not self.charm._patroni.is_replication_healthy(raft_encryption)
                     ):
                         logger.debug(
                             "Instance not yet back in the cluster."

tests/unit/test_cluster.py

@@ -267,6 +267,11 @@ def test_is_replication_healthy(peers_ips, patroni):
         ]
         assert not patroni.is_replication_healthy()
 
+        # Test ignoring errors in case of raft encryption.
+        _get.side_effect = None
+        _get.return_value.status_code = 503
+        assert patroni.is_replication_healthy(True)
+
 
 def test_is_member_isolated(peers_ips, patroni):
     with (