From 3bb61e0a75f12349be824c7460737aef917f49cd Mon Sep 17 00:00:00 2001
From: Dimitris Moraitis <dimitrios.moraitis@dell.com>
Date: Tue, 4 Jun 2024 14:09:00 +0300
Subject: [PATCH] Add support for overriding images of target containers

---
 examples/mounts-template.yaml |  2 ++
 operator/crd.yaml             |  5 +++++
 operator/op.py                | 33 +++++++++++++++++++++++++++++++--
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/examples/mounts-template.yaml b/examples/mounts-template.yaml
index dd9d3b8..d4f230e 100644
--- a/examples/mounts-template.yaml
+++ b/examples/mounts-template.yaml
@@ -4,3 +4,5 @@
     labels: ${LABELS}
     entrypoints:
       ${ENTRYPOINTS}
+    images:
+      ${IMAGES}
diff --git a/operator/crd.yaml b/operator/crd.yaml
index ad66b7a..fcc6df7 100644
--- a/operator/crd.yaml
+++ b/operator/crd.yaml
@@ -112,6 +112,11 @@ spec:
                       description: Override the entrypoint of specific containers in the target Deployment or Statefulset when mounting the code volume.
                       additionalProperties:
                         type: string
+                    images:
+                      type: object
+                      description: Override the image of specific containers in the target Deployment or Statefulset when mounting the code volume.
+                      additionalProperties:
+                        type: string
                     mounted:
                       type: boolean
                       description: Whether the volume should actually be mounted.
diff --git a/operator/op.py b/operator/op.py
index e4cd64b..9bc144a 100644
--- a/operator/op.py
+++ b/operator/op.py
@@ -78,12 +78,14 @@ def update_mounts(name, spec, namespace, logger, **kwargs):
     """This handler will idempotently update the volume mounts."""
     del kwargs
     logger.info("Will idempotently update volume mounts.")
+    # import ipdb;ipdb.set_trace()
     for (
         manifest,
         mounted,
         mount_path,
         sub_path,
         entrypoints,
+        images,
     ) in iter_mounts_and_manifests(namespace, spec["mounts"]):
         m_kind, m_name = manifest["kind"], manifest["metadata"]["name"]
 
@@ -120,11 +122,13 @@ def update_mounts(name, spec, namespace, logger, **kwargs):
                 sub_path=sub_path,
             )
             update_entrypoints(manifest=manifest, entrypoints=entrypoints)
+            update_images(manifest=manifest, images=images)
             kubectl_apply(namespace=namespace, manifest=manifest, logger=logger)
         else:
             if spec.get("mode") == "modify":
                 remove_mount(manifest=manifest, volume_name=name)
                 restore_entrypoints(manifest=manifest, entrypoints=entrypoints)
+                restore_images(manifest=manifest)
                 logger.info("Idempotently unmounting volume to %s:%s", m_kind, m_name)
                 kubectl_apply(namespace=namespace, manifest=manifest, logger=logger)
             elif kubectl_get(namespace=namespace, kind=m_kind, labels={"devenv": name}):
@@ -222,7 +226,7 @@ def kubectl_get(namespace: str, kind: str, labels: dict[str, str]) -> list[dict]
 def iter_mounts_and_manifests(namespace, mounts):
     for mount in mounts:
         assert isinstance(mount, dict), repr(mount)
-        for attr in ("kind", "labels", "mountPath", "mounted", "entrypoints"):
+        for attr in ("kind", "labels", "mountPath", "mounted", "entrypoints", "images"):
             assert attr in mount, mount
         if mount["kind"].lower() != "deployment":
             raise NotImplementedError("Only deployments are supported.")
@@ -235,6 +239,7 @@ def iter_mounts_and_manifests(namespace, mounts):
                 mount["mountPath"],
                 mount.get("subPath", ""),
                 mount["entrypoints"],
+                mount["images"],
             )
 
 
@@ -258,6 +263,7 @@ def add_mount(
         volumes.append(
             {"name": volume_name, "persistentVolumeClaim": {"claimName": pvc_name}}
         )
+        manifest["spec"]["template"]["spec"]["volumes"] = volumes
     # Configure volume mount.
     for container in manifest["spec"]["template"]["spec"]["containers"]:
         mounts = container.get("volumeMounts", [])
@@ -281,7 +287,7 @@ def add_mount(
 
 def remove_mount(*, manifest: dict, volume_name: str) -> None:
     # Remove volume.
-    volumes = manifest["spec"]["template"]["spec"]["volumes"]
+    volumes = manifest["spec"]["template"]["spec"].get("volumes", [])
     for i, volume in reversed(list(enumerate(volumes))):
         if volume["name"] == volume_name:
             volumes.pop(i)
@@ -311,3 +317,26 @@ def restore_entrypoints(*, manifest: dict, entrypoints: dict) -> None:
         if container.get("command"):
             del container["command"]
             del container["args"]
+
+
+def update_images(*, manifest: dict, images: dict) -> None:
+    for container in manifest["spec"]["template"]["spec"]["containers"]:
+        image = images.get(container["name"])
+        if image is None:
+            continue
+        manifest["metadata"]["annotations"][
+            f"devenv.dell.com/{container['name']}"
+        ] = container["image"]
+        container["image"] = image
+
+
+def restore_images(*, manifest: dict) -> None:
+    for container in manifest["spec"]["template"]["spec"]["containers"]:
+        image = manifest["metadata"]["annotations"].get(
+            f"devenv.dell.com/{container['name']}"
+        )
+        if image:
+            container["image"] = image
+            del manifest["metadata"]["annotations"][
+                f"devenv.dell.com/{container['name']}"
+            ]