Merge pull request #650 from yuvipanda/merge-properly

Let `kubespawner_override` merge instead of replace dictionaries

docs/source/changelog.md

@@ -9,6 +9,44 @@
 - Now `c.KubeSpawner.environment` values supports substitution, just like other config options [#642](https://github.com/jupyterhub/kubespawner/pull/642) ([@dolfinus](https://github.com/dolfinus)).
   For example, `{"MYVAR": "jupyterhub-{username}"}` will rendered as `{"MYVAR": "jupyterhub-sam"}` for a user named `sam`.
   This could break backward compatibility if environment variable value contains strings like `{something}` there `something` is a substitution variable unknown for the KubeSpawner. You should escape braces by using `{{something}}` syntax.
+- A breaking change could occur if `kubespawner_override` is used to
+  override a dictionary based configuration listed below. It will only
+  be a breaking change if there is a difference between merging and
+  replacing previous dictionary though.
+
+  There could be a difference between merging and replacing if for
+  example `common_labels` is configured ahead of time to be `{"a": "a"}`,
+  and then `kubespawner_override` is used to set it to `{"b": "b"}`. Then
+  previosly the result would be {"b": "b"} while now it will be `{"a": "a", "b": "b"}`.
+
+  Only `common_labels` has a default value provided by KubeSpawner or the
+  JupyterHub Helm chart, making this mainly an issue for users that
+  first have configured one of these values and then expect it to be
+  entirely replaced in `kubespawner_override`. As an example, someone may
+  have configured `node_selector` to be `{"a": "a"}` by default and then
+  expect it to become `{"b": "b"}` after using `kubespawner_override`, but
+  after this change it will end up being `{"a": "a", "b": "b"}`.
+
+  Audit your config for any of the following traitlets that could be overriden
+  in `kubespawner_override`, and make sure they behave as you expect.
+
+  ```
+  common_labels
+  environment
+  extra_annotations
+  extra_container_config
+  extra_labels
+  extra_pod_config
+  extra_resource_guarantees
+  extra_resource_limits
+  lifecycle_hooks
+  node_selector
+  storage_extra_annotations
+  storage_extra_labels
+  storage_selector
+  user_namespace_annotations
+  user_namespace_labels
+  ```
 
 ## 4.3
 

kubespawner/spawner.py

@@ -43,6 +43,7 @@
     make_service,
 )
 from .reflector import ResourceReflector
+from .utils import recursive_update
 
 
 class PodReflector(ResourceReflector):
@@ -1597,7 +1598,10 @@ def _validate_image_pull_secrets(self, proposal):
         - `kubespawner_override`: a dictionary with overrides to apply to the KubeSpawner
           settings. Each value can be either the final value to change or a callable that
           take the `KubeSpawner` instance as parameter and return the final value. This can
-          be further overridden by `profile_options`
+          be further overridden by 'profile_options'
+          If the traitlet being overriden is a *dictionary*, the dictionary
+          will be *recursively updated*, rather than overriden. If you want to
+          remove a key, set its value to `None`
         - `profile_options`: A dictionary of sub-options that allow users to further customize the
           selected profile. By default, these are rendered as a dropdown with the label
           provided by `display_name`. Items should have a unique key representing the customization,
@@ -1616,6 +1620,9 @@ def _validate_image_pull_secrets(self, proposal):
               and value can be either the final value or a callable that returns the final
               value when called with the spawner instance as the only parameter. The callable
               may be async.
+              If the traitlet being overriden is a *dictionary*, the dictionary
+              will be *recursively updated*, rather than overriden. If you want to
+              remove a key, set its value to `None`
         - `default`: (optional Bool) True if this is the default selected option
 
         kubespawner setting overrides work in the following manner, with items further in the
@@ -3029,7 +3036,15 @@ async def _load_profile(self, slug, selected_profile_user_options):
                 )
             else:
                 self.log.debug(".. overriding KubeSpawner value %s=%s", k, v)
-            setattr(self, k, v)
+
+            # If v is a dict, *merge* it with existing values, rather than completely
+            # resetting it. This allows *adding* things like environment variables rather
+            # than completely replacing them. If value is set to None, the key
+            # will be removed
+            if isinstance(v, dict) and isinstance(getattr(self, k), dict):
+                recursive_update(getattr(self, k), v)
+            else:
+                setattr(self, k, v)
 
         if profile.get('profile_options'):
             # each option specified here *must* have a value in our POST, as we
@@ -3073,7 +3088,15 @@ async def _load_profile(self, slug, selected_profile_user_options):
                         self.log.debug(
                             f'.. overriding traitlet {k}={v} for option {option_name}={chosen_option}'
                         )
-                    setattr(self, k, v)
+
+                    # If v is a dict, *merge* it with existing values, rather than completely
+                    # resetting it. This allows *adding* things like environment variables rather
+                    # than completely replacing them. If value is set to None, the key
+                    # will be removed
+                    if isinstance(v, dict) and isinstance(getattr(self, k), dict):
+                        recursive_update(getattr(self, k), v)
+                    else:
+                        setattr(self, k, v)
 
     # set of recognised user option keys
     # used for warning about ignoring unrecognised options

kubespawner/utils.py

@@ -202,3 +202,23 @@ def host_matching(host: str, wildcard: str) -> bool:
     # user.example.com !~ *.user.example.com
     # user.example.com !~ *.example
     return host_parts[1:] == wildcard_parts[1:]
+
+
+# From https://github.com/jupyter-server/jupyter_server/blob/fc0ac3236fdd92778ea765db6e8982212c8389ee/jupyter_server/config_manager.py#L14
+def recursive_update(target, new):
+    """
+    Recursively update one dictionary in-place using another.
+
+    None values will delete their keys.
+    """
+    for k, v in new.items():
+        if isinstance(v, dict):
+            if k not in target:
+                target[k] = {}
+            recursive_update(target[k], v)
+
+        elif v is None:
+            target.pop(k, None)
+
+        else:
+            target[k] = v

tests/test_spawner.py

@@ -703,6 +703,7 @@ async def test_get_pod_manifest_tolerates_mixed_input():
             'image': 'training/python:label',
             'cpu_limit': 1,
             'mem_limit': 512 * 1024 * 1024,
+            'environment': {'override': 'override-value'},
         },
     },
     {
@@ -714,6 +715,16 @@ async def test_get_pod_manifest_tolerates_mixed_input():
             'mem_limit': 8 * 1024 * 1024 * 1024,
         },
     },
+    {
+        'display_name': 'Training Env - R',
+        'slug': 'training-r',
+        'kubespawner_override': {
+            'image': 'training/r:label',
+            'cpu_limit': 1,
+            'mem_limit': 512 * 1024 * 1024,
+            'environment': {'override': 'override-value', "to-remove": None},
+        },
+    },
 ]
 
 
@@ -735,6 +746,28 @@ async def test_user_options_set_from_form():
         assert getattr(spawner, key) == value
 
 
+async def test_kubespawner_override():
+    spawner = KubeSpawner(_mock=True)
+    spawner.profile_list = _test_profiles
+    # Set a base environment
+    # to-remove will be removed because we set its value to None
+    # in the override
+    spawner.environment = {"existing": "existing-value", "to-remove": "does-it-matter"}
+    # render the form, select first option
+    await spawner.get_options_form()
+    spawner.user_options = spawner.options_from_form(
+        {'profile': [_test_profiles[2]['slug']]}
+    )
+    assert spawner.user_options == {
+        'profile': _test_profiles[2]['slug'],
+    }
+    await spawner.load_user_options()
+    assert spawner.environment == {
+        "existing": "existing-value",
+        "override": "override-value",
+    }
+
+
 async def test_user_options_api():
     spawner = KubeSpawner(_mock=True)
     spawner.profile_list = _test_profiles