Support a charm resource to override the installed snap

.licenserc.yaml

@@ -19,6 +19,7 @@ header:
       - 'charms/worker/k8s/lib/charms/k8s/**'
     paths-ignore:
       - 'charms/worker/k8s/lib/charms/**'
+      - 'tests/integration/data/*.tar.gz'
       - '.github/**'
       - '**/.gitkeep'
       - '**/*.cfg'

charms/worker/charmcraft.yaml

@@ -43,6 +43,7 @@ bases:
     - name: ubuntu
       channel: "24.04"
       architectures: [amd64]
+
 config:
   options:
     labels:
@@ -54,6 +55,22 @@ config:
 
         Note: Due to NodeRestriction, workers are limited to how they can label themselves
         https://kubernetes.io/docs/reference/access-authn-authz/admission-controllers/#noderestriction
+
+resources:
+  snap-installation:
+    type: file
+    filename: snap-installation.tar.gz
+    description: |
+      Override charm defined snap installation script
+
+      This charm is designed to operate with a specific revision of snaps, overriding 
+      with anything will indicate that the charm is running an unsupported configuration.
+
+      Content Options:
+      0-byte resource (Default) -- Use the charm defined snap installation script
+      ./snap-installation.yaml  -- Overrides the charm defined snap-installation.yaml
+      ./k8s_XXXX.snap           -- Overrides the charm with a specific snap file installed dangerously
+
 parts:
   charm:
     plugin: charm
@@ -74,6 +91,7 @@ parts:
 provides:
   cos-agent:
     interface: cos_agent
+
 requires:
   cluster:
     interface: k8s-cluster

charms/worker/k8s/charmcraft.yaml

@@ -135,6 +135,21 @@ config:
         only possible to increase the size of the IP range. It is not possible to
         change or shrink the address range after deployment.
 
+resources:
+  snap-installation:
+    type: file
+    filename: snap-installation.tar.gz
+    description: |
+      Override charm defined snap installation script
+
+      This charm is designed to operate with a specific revision of snaps, overriding 
+      with anything will indicate that the charm is running an unsupported configuration.
+
+      Content Options:
+      0-byte resource (Default) -- Use the charm defined snap installation script
+      ./snap-installation.yaml  -- Overrides the charm defined snap-installation.yaml
+      ./k8s_XXXX.snap           -- Overrides the charm with a specific snap file installed dangerously
+
 actions:
   get-kubeconfig:
     description: Retrieve Public Kubernetes cluster config, including credentials

charms/worker/k8s/src/charm.py

@@ -103,6 +103,38 @@ class NodeRemovedError(Exception):
     """Raised to prevent reconciliation of dying node."""
 
 
+class DynamicActiveStatus(status.ActiveStatus):
+    """An ActiveStatus class that can be updated.
+
+    Attributes:
+        message (str): explanation of the unit status
+        prefix  (str): Optional prefix to the unit status
+        postfix (str): Optional postfix to the unit status
+    """
+
+    def __init__(self):
+        """Initialise the DynamicActiveStatus."""
+        super().__init__("Ready")
+        self.prefix: str = ""
+        self.postfix: str = ""
+
+    @property
+    def message(self) -> str:
+        """Return the message for the status."""
+        pre = f"{self.prefix} :" if self.prefix else ""
+        post = f" ({self.postfix})" if self.postfix else ""
+        return f"{pre}{self._message}{post}"
+
+    @message.setter
+    def message(self, message: str):
+        """Set the message for the status.
+
+        Args:
+            message (str): explanation of the unit status
+        """
+        self._message = message
+
+
 class K8sCharm(ops.CharmBase):
     """A charm for managing a K8s cluster via the k8s snap.
 
@@ -126,7 +158,8 @@ def __init__(self, *args):
         xcp_relation = "external-cloud-provider" if self.is_control_plane else ""
         self.xcp = ExternalCloudProvider(self, xcp_relation)
         self.cos = COSIntegration(self)
-        self.reconciler = Reconciler(self, self._reconcile)
+        self.active_status = DynamicActiveStatus()
+        self.reconciler = Reconciler(self, self._reconcile, exit_status=self.active_status)
         self.distributor = TokenDistributor(self, self.get_node_name(), self.api_manager)
         self.collector = TokenCollector(self, self.get_node_name())
         self.labeller = LabelMaker(
@@ -247,7 +280,7 @@ def get_cloud_name(self) -> str:
     def _install_snaps(self):
         """Install snap packages."""
         status.add(ops.MaintenanceStatus("Ensuring snap installation"))
-        snap_management()
+        snap_management(self)
 
     @on_error(WaitingStatus("Waiting to apply snap requirements"), subprocess.CalledProcessError)
     def _apply_snap_requirements(self):
@@ -628,9 +661,12 @@ def _reconcile(self, event: ops.EventBase):
 
     def _update_status(self):
         """Check k8s snap status."""
-        if version := snap_version("k8s"):
+        version, overridden = snap_version("k8s")
+        if version:
             self.unit.set_workload_version(version)
 
+        self.active_status.postfix = "Snap Override Active" if overridden else ""
+
         if not self.get_cluster_name():
             status.add(ops.WaitingStatus("Node not Clustered"))
             return
@@ -743,7 +779,7 @@ def _on_update_status(self, _event: ops.UpdateStatusEvent):
             return
 
         try:
-            with status.context(self.unit):
+            with status.context(self.unit, exit_status=self.active_status):
                 self._update_status()
         except status.ReconcilerError:
             log.exception("Can't update_status")

charms/worker/k8s/src/snap.py

@@ -10,11 +10,14 @@
 
 import logging
 import re
+import shlex
+import shutil
 import subprocess
 from pathlib import Path
-from typing import List, Literal, Optional, Union
+from typing import List, Literal, Optional, Tuple, Union
 
 import charms.operator_libs_linux.v2.snap as snap_lib
+import ops
 import yaml
 from pydantic import BaseModel, Field, ValidationError, parse_obj_as, validator
 from typing_extensions import Annotated
@@ -31,7 +34,7 @@ class SnapFileArgument(BaseModel):
         name (str): The name of the snap after installed
         filename (Path): Path to the snap to locally install
         classic (bool): If it should be installed as a classic snap
-        dangerous (bool): If it should be installed as a dangerouse snap
+        dangerous (bool): If it should be installed as a dangerous snap
         devmode (bool): If it should be installed as with dev mode enabled
     """
 
@@ -91,25 +94,140 @@ def _validate_revision(cls, value: Union[str, int, None]) -> Optional[str]:
 ]
 
 
-def _parse_management_arguments() -> List[SnapArgument]:
+def _local_arch() -> str:
+    """Retrieve the local architecture.
+
+    Returns:
+        str: The architecture of this machine
+    """
+    dpkg_arch = ["dpkg", "--print-architecture"]
+    return subprocess.check_output(dpkg_arch).decode("UTF-8").strip()
+
+
+def _default_snap_installation() -> Path:
+    """Return the default snap_installation manifest.
+
+    Returns:
+        path to the default snap_installation manifest
+    """
+    return Path("templates/snap_installation.yaml")
+
+
+def _overridden_snap_installation() -> Path:
+    """Return the overridden snap_installation manifest.
+
+    Returns:
+        path to the overridden snap_installation manifest
+    """
+    return Path("./snap-installation/resource/snap_installation.yaml")
+
+
+def _normalize_paths(snap_installation):
+    """Normalize the paths in the snap_installation manifest.
+
+    Arguments:
+        snap_installation: The path to the snap_installation manifest
+    """
+    snap_installation = snap_installation.resolve()
+    content = yaml.safe_load(snap_installation.read_text(encoding="utf-8"))
+    updated = False
+    for arch, snaps in content.items():
+        for idx, snap in enumerate(snaps):
+            if snap.get("filename"):
+                resolved = (snap_installation.parent / snap["filename"]).resolve()
+                log.info("Resolving snap filename: %s to %s", snap["filename"], resolved)
+                content[arch][idx]["filename"] = str(resolved)
+                updated = True
+    if updated:
+        yaml.safe_dump(content, snap_installation.open(mode="w", encoding="utf-8"))
+
+
+def _select_snap_installation(charm: ops.CharmBase) -> Path:
+    """Select the snap_installation manifest.
+
+    Arguments:
+        charm: The charm instance necessary to check the unit resources
+
+    Returns:
+        path: The path to the snap_installation manifest
+
+    Raises:
+        SnapError: when the management issue cannot be resolved
+    """
+    try:
+        resource_path = charm.model.resources.fetch("snap-installation")
+    except (ops.ModelError, NameError):
+        log.error("Something went wrong when claiming 'snap-installation' resource.")
+        return _default_snap_installation()
+
+    resource_size = resource_path.stat().st_size
+    log.info("Resource path size: %d bytes", resource_size)
+    unpack_path = _overridden_snap_installation().parent
+    shutil.rmtree(unpack_path, ignore_errors=True)
+    if resource_size == 0:
+        log.info("Resource size is zero bytes. Use the charm defined snap installation script")
+        return _default_snap_installation()
+
+    # Unpack the snap-installation resource
+    unpack_path.mkdir(parents=True, exist_ok=True)
+    command = f"tar -xzvf {resource_path} -C {unpack_path} --no-same-owner"
+    try:
+        subprocess.check_call(shlex.split(command))
+    except subprocess.CalledProcessError as e:
+        log.error("Failed to extract 'snap-installation:'")
+        raise snap_lib.SnapError("Invalid snap-installation resource") from e
+
+    # Find the snap_installation manifest
+    snap_installation = unpack_path / "snap_installation.yaml"
+    if snap_installation.exists():
+        log.info("Found snap_installation manifest")
+        _normalize_paths(snap_installation)
+        return snap_installation
+
+    snap_path = list(unpack_path.glob("*.snap"))
+    if len(snap_path) == 1:
+        log.info("Found snap_installation snap: %s", snap_path[0])
+        arch = _local_arch()
+        manifest = {
+            arch: [
+                {
+                    "install-type": "file",
+                    "name": "k8s",
+                    "filename": str(snap_path[0]),
+                    "classic": True,
+                    "dangerous": True,
+                }
+            ]
+        }
+        yaml.safe_dump(manifest, snap_installation.open("w"))
+        return snap_installation
+
+    log.error("Failed to find a snap file in snap_installation resource")
+    raise snap_lib.SnapError("Failed to find snap_installation manifest")
+
+
+def _parse_management_arguments(charm: ops.CharmBase) -> List[SnapArgument]:
     """Parse snap management arguments.
 
+    Arguments:
+        charm: The charm instance necessary to check the unit resources
+
     Raises:
         SnapError: when the management issue cannot be resolved
 
     Returns:
         Parsed arguments list for the specific host architecture
     """
-    revision = Path("templates/snap_installation.yaml")
+    revision = _select_snap_installation(charm)
     if not revision.exists():
         raise snap_lib.SnapError(f"Failed to find file={revision}")
     try:
         body = yaml.safe_load(revision.read_text(encoding="utf-8"))
     except yaml.YAMLError as e:
         log.error("Failed to load file=%s, %s", revision, e)
         raise snap_lib.SnapError(f"Failed to load file={revision}")
-    dpkg_arch = ["dpkg", "--print-architecture"]
-    arch = subprocess.check_output(dpkg_arch).decode("UTF-8").strip()
+
+    arch = _local_arch()
 
     if not (isinstance(body, dict) and (arch_spec := body.get(arch))):
         log.warning("Failed to find revision for arch=%s", arch)
@@ -126,10 +244,14 @@ def _parse_management_arguments() -> List[SnapArgument]:
     return args
 
 
-def management():
-    """Manage snap installations on this machine."""
+def management(charm: ops.CharmBase) -> None:
+    """Manage snap installations on this machine.
+
+    Arguments:
+        charm: The charm instance
+    """
     cache = snap_lib.SnapCache()
-    for args in _parse_management_arguments():
+    for args in _parse_management_arguments(charm):
         which = cache[args.name]
         if isinstance(args, SnapFileArgument) and which.revision != "x1":
             snap_lib.install_local(**args.dict(exclude_none=True))
@@ -143,7 +265,7 @@ def management():
             which.ensure(**args.dict(exclude_none=True))
 
 
-def version(snap: str) -> Optional[str]:
+def version(snap: str) -> Tuple[Optional[str], bool]:
     """Retrieve the version of the installed snap package.
 
     Arguments:
@@ -153,15 +275,16 @@ def version(snap: str) -> Optional[str]:
         Optional[str]: The version of the installed snap package, or None if
         not available.
     """
+    overridden = _overridden_snap_installation().exists()
     try:
         result = subprocess.check_output(["/usr/bin/snap", "list", snap])
     except subprocess.CalledProcessError:
-        return None
+        return None, overridden
 
     output = result.decode().strip()
     match = re.search(r"(\d+\.\d+(?:\.\d+)?)", output)
     if match:
-        return match.group()
+        return match.group(), overridden
 
     log.info("Snap k8s not found or no version available.")
-    return None
+    return None, overridden