DPE-2919 update pod labels on preemptive switchover (#367)

* primary switchover update pod labels immediately

* shorter version

* fix pod labeling detail and removed duplicated method

* fix log and safe exceptin

Author: paulomach

src/k8s_helpers.py

@@ -113,21 +113,26 @@ def label_pod(self, role: str, pod_name: Optional[str] = None) -> None:
             pod_name: (optional) name of the pod to label, defaults to the current pod
         """
         try:
+            if not pod_name:
+                pod_name = self.pod_name
             pod = self.client.get(Pod, pod_name or self.pod_name, namespace=self.namespace)
 
             if not pod.metadata.labels:
                 pod.metadata.labels = {}
 
             if pod.metadata.labels.get("role") == role:
                 return
+            logger.debug(f"Patching {pod_name=} with {role=}")
 
             pod.metadata.labels["cluster-name"] = self.cluster_name
             pod.metadata.labels["role"] = role
-            self.client.patch(Pod, pod_name or self.pod_name, pod)
-            logger.info(f"Kubernetes pod label {role} created")
+            self.client.patch(Pod, pod_name, pod)
         except ApiError as e:
             if e.status.code == 404:
-                logger.warning(f"Kubernetes pod {pod_name} not found. Scaling in?")
+                logger.warning(f"Kubernetes {pod_name=} not found. Scaling in?")
+                return
+            if e.status.code == 409:
+                logger.warning(f"Kubernetes {pod_name=} changed. Labeling skipped")
                 return
             if e.status.code == 403:
                 logger.error("Kubernetes pod label creation failed: `juju trust` needed")

src/mysql_k8s_helpers.py

@@ -15,6 +15,7 @@
     MySQLClientError,
     MySQLConfigureMySQLUsersError,
     MySQLExecError,
+    MySQLGetClusterEndpointsError,
     MySQLStartMySQLDError,
     MySQLStopMySQLDError,
 )
@@ -47,7 +48,7 @@
     ROOT_SYSTEM_USER,
     XTRABACKUP_PLUGIN_DIR,
 )
-from k8s_helpers import KubernetesHelpers
+from k8s_helpers import KubernetesClientError, KubernetesHelpers
 from utils import any_memory_to_bytes
 
 logger = logging.getLogger(__name__)
@@ -613,6 +614,7 @@ def _execute_commands(
         user: Optional[str] = None,
         group: Optional[str] = None,
         env_extra: Optional[Dict] = None,
+        timeout: Optional[float] = None,
     ) -> Tuple[str, str]:
         """Execute commands on the server where MySQL is running."""
         try:
@@ -624,6 +626,7 @@ def _execute_commands(
                 user=user,
                 group=group,
                 environment=env_extra,
+                timeout=timeout,
             )
             stdout, stderr = process.wait_output()
             return (stdout, stderr or "")
@@ -820,3 +823,27 @@ def is_data_dir_initialised(self) -> bool:
             return expected_content <= content_set
         except ExecError:
             return False
+
+    def update_endpoints(self) -> None:
+        """Updates pod labels to reflect role of the unit."""
+        logger.debug("Updating pod labels")
+        try:
+            rw_endpoints, ro_endpoints, offline = self.get_cluster_endpoints(get_ips=False)
+
+            for endpoints, label in (
+                (rw_endpoints, "primary"),
+                (ro_endpoints, "replicas"),
+                (offline, "offline"),
+            ):
+                for pod in (p.split(".")[0] for p in endpoints.split(",")):
+                    if pod:
+                        self.k8s_helper.label_pod(label, pod)
+        except MySQLGetClusterEndpointsError:
+            logger.exception("Failed to get cluster endpoints")
+        except KubernetesClientError:
+            logger.exception("Can't update pod labels")
+
+    def set_cluster_primary(self, new_primary_address: str) -> None:
+        """Set the cluster primary and update pod labels."""
+        super().set_cluster_primary(new_primary_address)
+        self.update_endpoints()

src/relations/mysql_provider.py

@@ -6,7 +6,6 @@
 import logging
 import socket
 import typing
-from typing import List
 
 from charms.data_platform_libs.v0.data_interfaces import (
     DatabaseProvides,
@@ -16,7 +15,6 @@
     MySQLCreateApplicationDatabaseAndScopedUserError,
     MySQLDeleteUserError,
     MySQLDeleteUsersForRelationError,
-    MySQLGetClusterEndpointsError,
     MySQLGetMySQLVersionError,
     MySQLGrantPrivilegesToUserError,
     MySQLRemoveRouterFromMetadataError,
@@ -84,65 +82,6 @@ def _get_or_set_password(self, relation) -> str:
         self.database.update_relation_data(relation.id, {"password": password})
         return password
 
-    @staticmethod
-    def _endpoints_to_pod_list(endpoints: str) -> List[str]:
-        """Converts a comma separated list of endpoints to a list of pods."""
-        return [p.split(".")[0] for p in endpoints.split(",")]
-
-    def _update_endpoints(self) -> None:
-        """Updates pod labels to reflect role of the unit."""
-        logger.debug("Updating pod labels")
-        try:
-            rw_endpoints, ro_endpoints, offline = self.charm._mysql.get_cluster_endpoints(
-                get_ips=False
-            )
-
-            # rw pod labels
-            if rw_endpoints:
-                for pod in self._endpoints_to_pod_list(rw_endpoints):
-                    self.charm.k8s_helpers.label_pod("primary", pod)
-            # ro pod labels
-            if ro_endpoints:
-                for pod in self._endpoints_to_pod_list(ro_endpoints):
-                    self.charm.k8s_helpers.label_pod("replicas", pod)
-            # offline pod labels
-            if offline:
-                for pod in self._endpoints_to_pod_list(offline):
-                    self.charm.k8s_helpers.label_pod("offline", pod)
-        except MySQLGetClusterEndpointsError:
-            logger.exception("Failed to get cluster endpoints")
-        except KubernetesClientError:
-            logger.debug("Can't update pod labels")
-            self.charm.unit.status = BlockedStatus("Can't update pod labels")
-
-    def _update_pod_endpoint(self) -> None:
-        """Update pod label to reflect the role of the unit."""
-        logger.debug(f"Updating pod endpoint for {self.charm.unit.name}")
-
-        pod = self.charm.unit.name.replace("/", "-")
-
-        try:
-            cluster_status = self.charm._mysql.get_cluster_status()
-            if not cluster_status:
-                self.charm.k8s_helpers.label_pod("error", pod)
-                return
-
-            for hostname, properties in cluster_status["defaultreplicaset"]["topology"].items():
-                if hostname.split(".")[0] == pod:
-                    if properties["status"] != "online":
-                        label = "offline"
-                    elif properties["memberrole"] == "secondary":
-                        label = "replicas"
-                    elif properties["memberrole"] == "primary":
-                        label = "primary"
-                    else:
-                        label = "none"
-
-                    logger.debug(f"Labeling pod {pod} with label {label}")
-                    self.charm.k8s_helpers.label_pod(label, pod)
-        except KubernetesClientError:
-            logger.error("Error updating pod label. Traffic may not be properly routed.")
-
     # =============
     # Handlers
     # =============
@@ -171,7 +110,7 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
 
         try:
             # make sure pods are labeled before adding service
-            self._update_endpoints()
+            self.charm._mysql.update_endpoints()
 
             # create k8s services for endpoints
             self.charm.k8s_helpers.create_endpoint_services(["primary", "replicas"])
@@ -267,7 +206,7 @@ def _configure_endpoints(self, _) -> None:
             if relation.id not in relation_data:
                 continue
 
-            self._update_pod_endpoint()
+            self.charm._mysql.update_endpoints()
 
     def _on_update_status(self, _) -> None:
         """Handle the update status event.
@@ -285,7 +224,7 @@ def _on_update_status(self, _) -> None:
         ):
             return
 
-        self._update_pod_endpoint()
+        self.charm._mysql.update_endpoints()
 
     def _on_database_broken(self, event: RelationBrokenEvent) -> None:
         """Handle the removal of database relation.

tests/unit/test_database.py

@@ -51,7 +51,7 @@ def tearDown(self) -> None:
         self.patcher.stop()
 
     @patch("k8s_helpers.KubernetesHelpers.wait_service_ready")
-    @patch("relations.mysql_provider.MySQLProvider._update_endpoints")
+    @patch("mysql_k8s_helpers.MySQL.update_endpoints")
     @patch("k8s_helpers.KubernetesHelpers.create_endpoint_services")
     @patch("mysql_k8s_helpers.MySQL.get_mysql_version", return_value="8.0.29-0ubuntu0.20.04.3")
     @patch("mysql_k8s_helpers.MySQL.create_application_database_and_scoped_user")

tests/unit/test_mysql_k8s_helpers.py

@@ -3,7 +3,7 @@
 
 import json
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
 
 import tenacity
 from charms.mysql.v0.mysql import MySQLClientError, MySQLConfigureMySQLUsersError
@@ -336,3 +336,31 @@ def test_wait_until_unit_removed_from_cluster_exception(self, _get_cluster_statu
 
         with self.assertRaises(MySQLWaitUntilUnitRemovedFromClusterError):
             self.mysql._wait_until_unit_removed_from_cluster("mysql-0.mysql-endpoints")
+
+    @patch(
+        "mysql_k8s_helpers.MySQL.get_cluster_endpoints",
+        return_value=(
+            "mysql-0.mysql-endpoints",
+            "mysql-1.mysql-endpoints,mysql-2.mysql-endpoints",
+            "mysql-3.mysql-endpoints",
+        ),
+    )
+    def test_update_endpoints(self, _get_cluster_endpoints):
+        """Test the successful execution of update_endpoints."""
+        _label_pod = MagicMock()
+        _mock_k8s_helper = MagicMock()
+        _mock_k8s_helper.label_pod = _label_pod
+
+        self.mysql.k8s_helper = _mock_k8s_helper
+
+        calls = [
+            call("primary", "mysql-0"),
+            call("replicas", "mysql-1"),
+            call("replicas", "mysql-2"),
+            call("offline", "mysql-3"),
+        ]
+
+        self.mysql.update_endpoints()
+        _get_cluster_endpoints.assert_called_once()
+
+        _label_pod.assert_has_calls(calls)