Add upgrade UnitState enum (#233)

Ported from canonical/mysql-router-operator#134

Author: carlcsaposs-canonical

src/abstract_charm.py

@@ -254,7 +254,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             # Run before checking `self._upgrade.is_compatible` in case incompatible upgrade was
             # forced & completed on all units.
             self._upgrade.set_versions_in_app_databag()
-        if self._upgrade.unit_state == "restarting":  # Kubernetes only
+        if self._upgrade.unit_state is upgrade.UnitState.RESTARTING:  # Kubernetes only
             if not self._upgrade.is_compatible:
                 logger.info(
                     "Upgrade incompatible. If you accept potential *data loss* and *downtime*, you can continue with `resume-upgrade force=true`"
@@ -268,7 +268,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             if not self._upgrade.is_compatible:
                 self.set_status(event=event)
                 return
-            if self._upgrade.unit_state == "outdated":
+            if self._upgrade.unit_state is upgrade.UnitState.OUTDATED:
                 if self._upgrade.authorized:
                     self._upgrade.upgrade_unit(
                         event=event,
@@ -330,7 +330,7 @@ def reconcile(self, event=None) -> None:  # noqa: C901
             # Empty waiting status means we're waiting for database requires relation before
             # starting workload
             if not workload_.status or workload_.status == ops.WaitingStatus():
-                self._upgrade.unit_state = "healthy"
+                self._upgrade.unit_state = upgrade.UnitState.HEALTHY
             if self._unit_lifecycle.authorized_leader:
                 self._upgrade.reconcile_partition()
             self.set_status(event=event)

src/kubernetes_charm.py

@@ -286,7 +286,7 @@ def _on_stop(self, _) -> None:
         if not self._upgrade:
             logger.debug("Peer relation missing during stop event")
             return
-        self._upgrade.unit_state = "restarting"
+        self._upgrade.unit_state = upgrade.UnitState.RESTARTING
 
 
 if __name__ == "__main__":

src/kubernetes_upgrade.py

@@ -159,7 +159,7 @@ def _app_workload_container_version(self) -> str:
         )
         return stateful_set.status.updateRevision
 
-    def reconcile_partition(self, *, action_event: ops.ActionEvent = None) -> None:
+    def reconcile_partition(self, *, action_event: ops.ActionEvent = None) -> None:  # noqa: C901
         """If ready, lower partition to upgrade next unit.
 
         If upgrade is not in progress, set partition to 0. (If a unit receives a stop event, it may
@@ -182,8 +182,11 @@ def determine_partition() -> int:
             logger.debug(f"{self._peer_relation.data=}")
             for upgrade_order_index, unit in enumerate(units):
                 # Note: upgrade_order_index != unit number
+                state = self._peer_relation.data[unit].get("state")
+                if state:
+                    state = upgrade.UnitState(state)
                 if (
-                    not force and self._peer_relation.data[unit].get("state") != "healthy"
+                    not force and state is not upgrade.UnitState.HEALTHY
                 ) or self._unit_workload_container_versions[
                     unit.name
                 ] != self._app_workload_container_version:
@@ -227,7 +230,7 @@ def determine_partition() -> int:
                 # not which units upgrade. Kubernetes may not upgrade a unit even if the partition
                 # allows it (e.g. if the charm container of a higher unit is not ready). This is
                 # also applicable `if not force`, but is unlikely to happen since all units are
-                # "healthy" `if not force`.
+                # healthy `if not force`.
                 message = f"Attempting to upgrade unit {self._partition}"
             else:
                 message = f"Upgrade resumed. Unit {self._partition} is upgrading next"

src/upgrade.py

@@ -9,6 +9,7 @@
 
 import abc
 import copy
+import enum
 import json
 import logging
 import pathlib
@@ -34,6 +35,15 @@ class PeerRelationNotReady(Exception):
     """Upgrade peer relation not available (to this unit)"""
 
 
+class UnitState(str, enum.Enum):
+    """Unit upgrade state"""
+
+    HEALTHY = "healthy"
+    RESTARTING = "restarting"  # Kubernetes only
+    UPGRADING = "upgrading"  # Machines only
+    OUTDATED = "outdated"  # Machines only
+
+
 class Upgrade(abc.ABC):
     """In-place upgrades"""
 
@@ -55,13 +65,14 @@ def __init__(self, charm_: ops.CharmBase) -> None:
             self._current_versions[version] = pathlib.Path(file_name).read_text().strip()
 
     @property
-    def unit_state(self) -> typing.Optional[str]:
+    def unit_state(self) -> typing.Optional[UnitState]:
         """Unit upgrade state"""
-        return self._unit_databag.get("state")
+        if state := self._unit_databag.get("state"):
+            return UnitState(state)
 
     @unit_state.setter
-    def unit_state(self, value: str) -> None:
-        self._unit_databag["state"] = value
+    def unit_state(self, value: UnitState) -> None:
+        self._unit_databag["state"] = value.value
 
     @property
     def is_compatible(self) -> bool: