Improve Docstrings and Code Logic

charms/worker/k8s/src/charm.py

@@ -125,6 +125,7 @@ class K8sCharm(ops.CharmBase):
         is_control_plane: true if this is a control-plane unit
         lead_control_plane: true if this is a control-plane unit and its the leader
         is_upgrade_granted: true if the upgrade has been granted
+        datastore: the datastore used for Kubernetes
     """
 
     _stored = ops.StoredState()
@@ -144,7 +145,7 @@ def __init__(self, *args):
         self.cluster_inspector = ClusterInspector(kubeconfig_path=self._internal_kubeconfig)
         self.upgrade = K8sUpgrade(
             self,
-            node_manager=self.cluster_inspector,
+            cluster_inspector=self.cluster_inspector,
             relation_name="upgrade",
             substrate="vm",
             dependency_model=K8sDependenciesModel(**DEPENDENCIES),
@@ -230,6 +231,11 @@ def is_worker(self) -> bool:
         """Returns true if the unit is a worker."""
         return self.meta.name == "k8s-worker"
 
+    @property
+    def datastore(self) -> str:
+        """Return the datastore type."""
+        return str(self.config.get("bootstrap-datastore"))
+
     def get_worker_versions(self) -> Dict[str, List[ops.Unit]]:
         """Get the versions of the worker units.
 

charms/worker/k8s/src/literals.py

@@ -11,12 +11,19 @@
     "k8sd",
 ]
 
+K8S_DQLITE_SERVICE = "k8s-dqlite"
+
 K8S_CONTROL_PLANE_SERVICES = [
     "kube-apiserver",
+    K8S_DQLITE_SERVICE,
     "kube-controller-manager",
     "kube-scheduler",
 ]
 
+K8S_WORKER_SERVICES = [
+    "k8s-apiserver-proxy",
+]
+
 DEPENDENCIES = {
     # NOTE: Update the dependencies for the k8s-charm before releasing.
     "k8s_charm": {

charms/worker/k8s/src/protocols.py

@@ -23,6 +23,8 @@ class K8sCharmProtocol(ops.CharmBase):
         is_upgrade_granted (bool): Whether the upgrade is granted.
         lead_control_plane (bool): Whether the charm is the lead control plane.
         is_control_plane (bool): Whether the charm is a control plane.
+        is_worker (bool): Whether the charm is a worker.
+        datastore (str): The datastore for Kubernetes.
     """
 
     api_manager: K8sdAPIManager
@@ -32,6 +34,8 @@ class K8sCharmProtocol(ops.CharmBase):
     is_upgrade_granted: bool
     lead_control_plane: bool
     is_control_plane: bool
+    is_worker: bool
+    datastore: str
 
     def get_cluster_name(self) -> str:
         """Get the cluster name.

charms/worker/k8s/src/snap.py

@@ -273,7 +273,7 @@ def management(charm: K8sCharmProtocol) -> None:
     cache = snap_lib.SnapCache()
     for args in _parse_management_arguments(charm):
         which = cache[args.name]
-        if block_refresh(which, args) and not charm.is_upgrade_granted:
+        if block_refresh(which, args, charm.is_upgrade_granted):
             continue
         install_args = args.dict(exclude_none=True)
         if isinstance(args, SnapFileArgument) and which.revision != "x1":
@@ -288,12 +288,13 @@ def management(charm: K8sCharmProtocol) -> None:
             which.ensure(**install_args)
 
 
-def block_refresh(which: snap_lib.Snap, args: SnapArgument) -> bool:
+def block_refresh(which: snap_lib.Snap, args: SnapArgument, upgrade_granted: bool = False) -> bool:
     """Block snap refreshes if the snap is in a specific state.
 
     Arguments:
         which: The snap to check
         args: The snap arguments
+        upgrade_granted: If the upgrade is granted
 
     Returns:
         bool: True if the snap should be blocked from refreshing
@@ -304,6 +305,9 @@ def block_refresh(which: snap_lib.Snap, args: SnapArgument) -> bool:
     if _overridden_snap_installation().exists():
         log.info("Allowing %s snap refresh due to snap installation override", args.name)
         return False
+    if upgrade_granted:
+        log.info("Allowing %s snap refresh due to upgrade-granted", args.name)
+        return False
     if isinstance(args, SnapStoreArgument) and args.revision:
         if block := which.revision != args.revision:
             log.info("Blocking %s snap refresh to revision=%s", args.name, args.revision)

charms/worker/k8s/src/upgrade.py

@@ -18,7 +18,13 @@
 )
 from charms.operator_libs_linux.v2.snap import SnapError
 from inspector import ClusterInspector
-from literals import K8S_COMMON_SERVICES, K8S_CONTROL_PLANE_SERVICES, SNAP_NAME
+from literals import (
+    K8S_COMMON_SERVICES,
+    K8S_CONTROL_PLANE_SERVICES,
+    K8S_DQLITE_SERVICE,
+    K8S_WORKER_SERVICES,
+    SNAP_NAME,
+)
 from protocols import K8sCharmProtocol
 from pydantic import BaseModel
 from snap import management as snap_management
@@ -43,17 +49,17 @@ class K8sDependenciesModel(BaseModel):
 class K8sUpgrade(DataUpgrade):
     """A helper class for upgrading the k8s and k8s-worker charms."""
 
-    def __init__(self, charm: K8sCharmProtocol, node_manager: ClusterInspector, **kwargs):
+    def __init__(self, charm: K8sCharmProtocol, cluster_inspector: ClusterInspector, **kwargs):
         """Initialize the K8sUpgrade.
 
         Args:
             charm: The charm instance.
-            node_manager: The ClusterInspector instance.
+            cluster_inspector: The ClusterInspector instance.
             kwargs: Additional keyword arguments.
         """
         super().__init__(charm, **kwargs)
         self.charm = charm
-        self.node_manager = node_manager
+        self.cluster_inspector = cluster_inspector
 
     def set_upgrade_status(self, event: ops.UpdateStatusEvent) -> None:
         """Set the Juju upgrade status.
@@ -84,10 +90,10 @@ def pre_upgrade_check(self) -> None:
             ClusterNotReadyError: If the cluster is not ready for an upgrade.
         """
         try:
-            nodes = self.node_manager.get_nodes(
+            nodes = self.cluster_inspector.get_nodes(
                 labels={"juju-charm": "k8s-worker" if self.charm.is_worker else "k8s"}
             )
-            failing_pods = self.node_manager.verify_pods_running(["kube-system"])
+            failing_pods = self.cluster_inspector.verify_pods_running(["kube-system"])
         except ClusterInspector.ClusterInspectorError as e:
             raise ClusterNotReadyError(
                 message="Cluster is not ready for an upgrade",
@@ -113,20 +119,23 @@ def pre_upgrade_check(self) -> None:
                 resolution="Check the logs for the failing pods.",
             )
 
-    def _verify_worker_version(self) -> bool:
-        """Verify the worker version.
+    def _verify_worker_versions(self) -> bool:
+        """Verify that the k8s-worker charm versions meet the requirements.
+
+        This method verifies that all applications related to the cluster relation
+        satisfy the requirements of the k8s-worker charm.
 
         Returns:
-            True if all the worker versions meet the requirements, otherwise False.
+            bool: True if all worker versions meet the requirements, False otherwise.
         """
-        worker_version = self.charm.get_worker_versions()
-        if not worker_version:
+        worker_versions = self.charm.get_worker_versions()
+        if not worker_versions:
             return True
         dependency_model: DependencyModel = getattr(self.dependency_model, "k8s_service")
 
         incompatible = {
             version: units
-            for version, units in worker_version.items()
+            for version, units in worker_versions.items()
             if not verify_requirements(
                 version=version, requirement=dependency_model.dependencies["k8s-worker"]
             )
@@ -149,12 +158,12 @@ def _perform_upgrade(self, services: List[str]) -> None:
         Args:
             services: The services to stop and start during the upgrade.
         """
-        status.add(ops.MaintenanceStatus("Stopping k8s Services."))
-        stop(SNAP_NAME, services=services)
+        status.add(ops.MaintenanceStatus("Stopping the K8s services"))
+        stop(SNAP_NAME, services)
         status.add(ops.MaintenanceStatus("Upgrading the k8s snap."))
         snap_management(self.charm)
-        status.add(ops.MaintenanceStatus("Restarting k8s Services."))
-        start(SNAP_NAME, services=services)
+        status.add(ops.MaintenanceStatus("Starting the K8s services"))
+        start(SNAP_NAME, services)
 
     def _on_upgrade_granted(self, event: UpgradeGrantedEvent) -> None:
         """Handle the upgrade granted event.
@@ -180,7 +189,7 @@ def _upgrade(self, event: Union[ops.EventBase, ops.HookEvent]) -> None:
         status.add(ops.MaintenanceStatus("Upgrading the charm."))
 
         if self.charm.lead_control_plane:
-            if not self._verify_worker_version():
+            if not self._verify_worker_versions():
                 self.set_unit_failed(
                     cause="The k8s worker charm version does not meet the requirements."
                 )
@@ -192,9 +201,12 @@ def _upgrade(self, event: Union[ops.EventBase, ops.HookEvent]) -> None:
         services = (
             K8S_CONTROL_PLANE_SERVICES + K8S_COMMON_SERVICES
             if self.charm.is_control_plane
-            else K8S_COMMON_SERVICES
+            else K8S_COMMON_SERVICES + K8S_WORKER_SERVICES
         )
 
+        if K8S_DQLITE_SERVICE in services and self.charm.datastore == "dqlite":
+            services.remove(K8S_DQLITE_SERVICE)
+
         try:
             self._perform_upgrade(services=services)
             self.set_unit_completed()

charms/worker/k8s/tests/unit/test_upgrade.py

@@ -23,7 +23,7 @@ def setUp(self):
         self.node_manager = MagicMock(spec=ClusterInspector)
         self.upgrade = K8sUpgrade(
             self.charm,
-            node_manager=self.node_manager,
+            cluster_inspector=self.node_manager,
             relation_name="upgrade",
             substrate="vm",
             dependency_model=K8sDependenciesModel(
@@ -37,7 +37,7 @@ def setUp(self):
                     "k8s_service": {
                         "dependencies": {"k8s-worker": "^1.30, < 1.32"},
                         "name": "k8s",
-                        "upgrade_supported": ">=0.8",
+                        "upgrade_supported": "^1.30, < 1.32",
                         "version": "1.31.1",
                     },
                 }
@@ -121,14 +121,26 @@ def test_build_upgrade_stack_with_relation(self):
         self.assertEqual(sorted(result), [0, 1, 2])
         self.charm.model.get_relation.assert_called_once_with("cluster")
 
-    def test_verify_worker_version_compatible(self):
-        """Test _verify_worker_version returns True when worker version is compatible."""
+    def test_verify_worker_versions_compatible(self):
+        """Test _verify_worker_versions returns True when worker versions is compatible."""
         unit_1 = MagicMock(spec=ops.Unit)
         unit_1.name = "k8s-worker/0"
         unit_2 = MagicMock(spec=ops.Unit)
         unit_2.name = "k8s-worker/1"
-        self.charm.get_worker_versions.return_value = {"1.32.0": [unit_1], "1.31.5": [unit_2]}
+        self.charm.get_worker_versions.return_value = {"1.31.0": [unit_1], "1.31.5": [unit_2]}
 
-        result = self.upgrade._verify_worker_version()
+        result = self.upgrade._verify_worker_versions()
 
-        self.assertTrue(not result)
+        self.assertTrue(result)
+
+    def test_verify_worker_versions_incompatible(self):
+        """Test _verify_worker_versions returns False when worker versions is incompatible."""
+        unit_1 = MagicMock(spec=ops.Unit)
+        unit_1.name = "k8s-worker/0"
+        unit_2 = MagicMock(spec=ops.Unit)
+        unit_2.name = "k8s-worker/1"
+        self.charm.get_worker_versions.return_value = {"1.32.0": [unit_1], "1.33.0": [unit_2]}
+
+        result = self.upgrade._verify_worker_versions()
+
+        self.assertFalse(result)