From ea2cc43f00f81caa60e0a76aa40ca3084d87e883 Mon Sep 17 00:00:00 2001
From: Mike Barry <github@barry.io>
Date: Tue, 7 Nov 2023 10:37:09 -0500
Subject: [PATCH] Templatize the notebook name

---
 kubespawner/objects.py |  6 +++++-
 kubespawner/spawner.py | 29 ++++++++++++++++++++++++++++-
 tests/test_spawner.py  |  4 ++++
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/kubespawner/objects.py b/kubespawner/objects.py
index 9df8bf0b..dce5c443 100644
--- a/kubespawner/objects.py
+++ b/kubespawner/objects.py
@@ -122,6 +122,7 @@ def make_pod(
     ssl_secret_name=None,
     ssl_secret_mount_path=None,
     logger=None,
+    notebook_container_name='notebook',
 ):
     """
     Make a k8s pod specification for running a user notebook.
@@ -333,6 +334,9 @@ def make_pod(
 
     ssl_secret_mount_path:
         Specifies the name of the ssl secret mount path for the pod
+
+    notebook_container_name:
+        The name of the notebook container in the pod
     """
 
     pod = V1Pod()
@@ -529,7 +533,7 @@ def _get_env_var_deps(env):
             break
 
     notebook_container = V1Container(
-        name='notebook',
+        name=notebook_container_name,
         image=image,
         working_dir=working_dir,
         ports=[V1ContainerPort(name='notebook-port', container_port=port)],
diff --git a/kubespawner/spawner.py b/kubespawner/spawner.py
index 66e66dc7..a8713ba5 100644
--- a/kubespawner/spawner.py
+++ b/kubespawner/spawner.py
@@ -193,6 +193,9 @@ def __init__(self, *args, **kwargs):
         self.secret_name = self._expand_user_properties(self.secret_name_template)
 
         self.pvc_name = self._expand_user_properties(self.pvc_name_template)
+        self.notebook_container_name = self._expand_user_properties(
+            self.notebook_container_name_template
+        )
         if self.working_dir:
             self.working_dir = self._expand_user_properties(self.working_dir)
         if self.port == 0:
@@ -760,6 +763,22 @@ def _deprecated_changed(self, change):
         """,
     )
 
+    notebook_container_name_template = Unicode(
+        'notebook',
+        config=True,
+        help="""
+        Template to use to form the name of the notebook container in the pod.
+
+        `{username}`, `{userid}`, `{servername}`, `{hubnamespace}`,
+        `{unescaped_username}`, and `{unescaped_servername}` will be expanded if
+        found within strings of this configuration. The username and servername
+        come escaped to follow the `DNS label standard
+        <https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-label-names>`__.
+
+        Trailing `-` characters are stripped.
+        """,
+    )
+
     @validate('image_pull_secrets')
     def _validate_image_pull_secrets(self, proposal):
         if type(proposal['value']) == str:
@@ -1859,6 +1878,9 @@ def _build_pod_labels(self, extra_labels):
                 'hub.jupyter.org/servername': escapism.escape(
                     self.name, safe=self.safe_chars, escape_char='-'
                 ).lower(),
+                # we put the container name in a label so if the template is updated we don't
+                # lose track of which pods are no longer running and leave stranded pods.
+                'hub.jupyter.org/notebook_container_name': self.notebook_container_name,
             }
         )
         return labels
@@ -2031,6 +2053,7 @@ async def get_pod_manifest(self):
             ssl_secret_name=self.secret_name if self.internal_ssl else None,
             ssl_secret_mount_path=self.secret_mount_path,
             logger=self.log,
+            notebook_container_name=self.notebook_container_name,
         )
 
     def get_secret_manifest(self, owner_reference):
@@ -2209,7 +2232,11 @@ async def poll(self):
                 return 1
             for c in ctr_stat:
                 # return exit code if notebook container has terminated
-                if c["name"] == 'notebook':
+                notebook_container_name = pod["metadata"]["labels"].get(
+                    "hub.jupyter.org/notebook_container_name", "notebook"
+                )
+
+                if c["name"] == notebook_container_name:
                     if "terminated" in c["state"]:
                         # call self.stop to delete the pod
                         if self.delete_stopped_pods:
diff --git a/tests/test_spawner.py b/tests/test_spawner.py
index 9daa7cdc..08599f86 100644
--- a/tests/test_spawner.py
+++ b/tests/test_spawner.py
@@ -1618,6 +1618,10 @@ async def test_variable_expansion(ssl_app):
             "configured_value": {"schedulerName": "extra-pod-config-{username}"},
             "findable_in": ["pod"],
         },
+        "notebook_container_name_template": {
+            "configured_value": "notebook-container-name-template-{username}",
+            "findable_in": ["pod"],
+        },
     }
 
     c = Config()