From c0d9e407d3c2b502d59a86a90ec8427a541c0da3 Mon Sep 17 00:00:00 2001
From: Dragomir Penev <dragomir.penev@canonical.com>
Date: Tue, 9 Sep 2025 15:30:21 +0300
Subject: [PATCH] Add back raft encryption check

---
 src/cluster.py             |  8 +++++++-
 src/upgrade.py             | 10 +++++++++-
 tests/unit/test_cluster.py |  5 +++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/cluster.py b/src/cluster.py
index b3212b6b7e..3e09123480 100644
--- a/src/cluster.py
+++ b/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}"
diff --git a/src/upgrade.py b/src/upgrade.py
index 7728af28b0..ed430e0cf4 100644
--- a/src/upgrade.py
+++ b/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."
diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
index 3ee9ba4bf7..e46dde3c31 100644
--- a/tests/unit/test_cluster.py
+++ b/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 (