Skip to content

Commit 5a3092f

Browse files
committed
Revert changes to drop wildcard support and send client_envs in kspec
1 parent 269e4c2 commit 5a3092f

File tree

5 files changed

+58
-92
lines changed

5 files changed

+58
-92
lines changed

enterprise_gateway/services/kernels/handlers.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,16 +54,24 @@ async def post(self):
5454
# Start with the PATH from the current env. Do not provide the entire environment
5555
# which might contain server secrets that should not be passed to kernels.
5656
env = {"PATH": os.getenv("PATH", "")}
57-
# Whitelist environment variables from current process environment
57+
# Transfer inherited environment variables from current process
5858
env.update(
5959
{key: value for key, value in os.environ.items() if key in self.inherited_envs}
6060
)
61+
# Allow all KERNEL_* envs and those specified in client_envs and set from client. If this EG
62+
# instance is configured to accept all envs in the payload (i.e., client_envs == '*'), go ahead
63+
# and add those keys to the "working" allowed_envs list, otherwise, just transfer the configured envs.
64+
allowed_envs: List[str]
65+
if self.client_envs == ["*"]:
66+
allowed_envs = model["env"].keys()
67+
else:
68+
allowed_envs = self.client_envs
6169
# Allow KERNEL_* args and those allowed by configuration.
6270
env.update(
6371
{
6472
key: value
6573
for key, value in model["env"].items()
66-
if key.startswith("KERNEL_") or key in self.client_envs
74+
if key.startswith("KERNEL_") or key in allowed_envs
6775
}
6876
)
6977

enterprise_gateway/services/kernels/remotemanager.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -456,7 +456,7 @@ async def start_kernel(self, **kwargs):
456456

457457
def _capture_user_overrides(self, **kwargs):
458458
"""
459-
Make a copy of any whitelist or KERNEL_ env values provided by user. These will be injected
459+
Make a copy of any allowed or KERNEL_ env values provided by user. These will be injected
460460
back into the env after the kernelspec env has been applied. This enables defaulting behavior
461461
of the kernelspec env stanza that would have otherwise overridden the user-provided values.
462462
"""

enterprise_gateway/services/kernelspecs/kernelspec_cache.py

Lines changed: 2 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
"""Cache handling for kernel specs."""
44

55
import os
6-
from typing import Dict, Optional, Set, Union
6+
from typing import Dict, Optional, Union
77

88
from jupyter_client.kernelspec import KernelSpec
99
from jupyter_server.utils import ensure_async
@@ -135,36 +135,8 @@ def put_item(self, kernel_name: str, cache_item: Union[KernelSpec, CacheItemType
135135
If it determines the cache entry corresponds to a currently unwatched directory,
136136
that directory will be added to list of observed directories and scheduled accordingly.
137137
"""
138-
self.log.info(f"KernelSpecCache: adding/updating kernelspec: {kernel_name}")
139-
# Irrespective of cache enablement, add/update the 'metadata.client_envs' entry
140-
# with the set of configured values. If the stanza already exists in the kernelspec
141-
# update with the union of it and those values configured via `EnterpriseGatewayApp'.
142-
# We apply this logic here so that its only performed once for cached values or on
143-
# every retrieval when caching is not enabled.
144-
# Note: We only need to do this if we have a KernelSpec instance, since CacheItemType
145-
# instances will have already been modified.
146-
147-
# Create a set from the configured value, update it with the (potential) value
148-
# in the kernelspec, and apply the changes back to the kernelspec.
149-
150-
client_envs: Set[str] = set(self.parent.client_envs)
151-
kspec_client_envs: Set[str]
152-
if type(cache_item) is KernelSpec:
153-
kspec: KernelSpec = cache_item
154-
kspec_client_envs = set(kspec.metadata.get("client_envs", []))
155-
else:
156-
kspec_client_envs = set(cache_item["spec"].get("metadata", {}).get("client_envs", []))
157-
158-
client_envs.update(kspec_client_envs)
159-
if type(cache_item) is KernelSpec:
160-
kspec: KernelSpec = cache_item
161-
kspec.metadata["client_envs"] = list(client_envs)
162-
else:
163-
if "metadata" not in cache_item["spec"]:
164-
cache_item["spec"]["metadata"] = {}
165-
cache_item["spec"]["metadata"]["client_envs"] = list(client_envs)
166-
167138
if self.cache_enabled:
139+
self.log.info(f"KernelSpecCache: adding/updating kernelspec: {kernel_name}")
168140
if type(cache_item) is KernelSpec:
169141
cache_item = KernelSpecCache.kernel_spec_to_cache_item(cache_item)
170142

enterprise_gateway/tests/test_handlers.py

Lines changed: 41 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -644,44 +644,44 @@ def test_base_url(self):
644644
self.assertEqual(response.code, 200)
645645

646646

647-
# class TestWildcardEnvs(TestHandlers):
648-
# """Base class for jupyter-websocket mode tests that spawn kernels."""
649-
#
650-
# def setup_app(self):
651-
# """Configure JUPYTER_PATH so that we can use local kernelspec files for testing."""
652-
# super().setup_app()
653-
# # overwrite env_whitelist
654-
# self.app.env_whitelist = ["*"]
655-
#
656-
# @gen_test
657-
# def test_kernel_wildcard_env(self):
658-
# """Kernel should start with environment vars defined in the request."""
659-
# # Note: Since env_whitelist == '*', all values should be present.
660-
# kernel_body = json.dumps(
661-
# {
662-
# "name": "python",
663-
# "env": {
664-
# "KERNEL_FOO": "kernel-foo-value",
665-
# "OTHER_VAR1": "other-var1-value",
666-
# "OTHER_VAR2": "other-var2-value",
667-
# "TEST_VAR": "test-var-value",
668-
# },
669-
# }
670-
# )
671-
# ws = yield self.spawn_kernel(kernel_body)
672-
# req = self.execute_request(
673-
# "import os; "
674-
# 'print(os.getenv("KERNEL_FOO"), '
675-
# 'os.getenv("OTHER_VAR1"), '
676-
# 'os.getenv("OTHER_VAR2"), '
677-
# 'os.getenv("TEST_VAR"))'
678-
# )
679-
# ws.write_message(json_encode(req))
680-
# content = yield self.await_stream(ws)
681-
# self.assertEqual(content["name"], "stdout")
682-
# self.assertIn("kernel-foo-value", content["text"])
683-
# self.assertIn("other-var1-value", content["text"])
684-
# self.assertIn("other-var2-value", content["text"])
685-
# self.assertIn("test-var-value", content["text"])
686-
#
687-
# ws.close()
647+
class TestWildcardEnvs(TestHandlers):
648+
"""Base class for jupyter-websocket mode tests that spawn kernels."""
649+
650+
def setup_app(self):
651+
"""Configure JUPYTER_PATH so that we can use local kernelspec files for testing."""
652+
super().setup_app()
653+
# overwrite env_whitelist
654+
self.app.env_whitelist = ["*"]
655+
656+
@gen_test
657+
def test_kernel_wildcard_env(self):
658+
"""Kernel should start with environment vars defined in the request."""
659+
# Note: Since env_whitelist == '*', all values should be present.
660+
kernel_body = json.dumps(
661+
{
662+
"name": "python",
663+
"env": {
664+
"KERNEL_FOO": "kernel-foo-value",
665+
"OTHER_VAR1": "other-var1-value",
666+
"OTHER_VAR2": "other-var2-value",
667+
"TEST_VAR": "test-var-value",
668+
},
669+
}
670+
)
671+
ws = yield self.spawn_kernel(kernel_body)
672+
req = self.execute_request(
673+
"import os; "
674+
'print(os.getenv("KERNEL_FOO"), '
675+
'os.getenv("OTHER_VAR1"), '
676+
'os.getenv("OTHER_VAR2"), '
677+
'os.getenv("TEST_VAR"))'
678+
)
679+
ws.write_message(json_encode(req))
680+
content = yield self.await_stream(ws)
681+
self.assertEqual(content["name"], "stdout")
682+
self.assertIn("kernel-foo-value", content["text"])
683+
self.assertIn("other-var1-value", content["text"])
684+
self.assertIn("other-var2-value", content["text"])
685+
self.assertIn("test-var-value", content["text"])
686+
687+
ws.close()

enterprise_gateway/tests/test_kernelspec_cache.py

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import pytest
1313
from jupyter_client.kernelspec import KernelSpecManager, NoSuchKernel
1414

15-
from enterprise_gateway.enterprisegatewayapp import EnterpriseGatewayApp
1615
from enterprise_gateway.services.kernelspecs import KernelSpecCache
1716

1817

@@ -61,16 +60,10 @@ def environ(
6160

6261
# END - Remove once transition to jupyter_server occurs
6362

64-
GLOBAL_CLIENT_ENVS = ["TEST_VAR", "OTHER_VAR1", "OTHER_VAR2"]
65-
KSPEC_CLIENT_ENVS = ["TEST_VAR", "OTHER_VAR1", "OTHER_VAR3", "OTHER_VAR4"]
66-
UNION_CLIENT_ENVS = ["TEST_VAR", "OTHER_VAR1", "OTHER_VAR2", "OTHER_VAR3", "OTHER_VAR4"]
6763

6864
kernelspec_json = {
6965
"argv": ["cat", "{connection_file}"],
7066
"display_name": "Test kernel: {kernel_name}",
71-
"metadata": {
72-
"client_envs": KSPEC_CLIENT_ENVS,
73-
}
7467
}
7568

7669

@@ -110,22 +103,17 @@ def setup_kernelspecs(environ, kernelspec_location):
110103

111104
@pytest.fixture
112105
def kernel_spec_manager(environ, setup_kernelspecs):
113-
app = EnterpriseGatewayApp()
114-
app.client_envs = ["TEST_VAR", "OTHER_VAR1", "OTHER_VAR2"]
115-
116-
yield KernelSpecManager(ensure_native_kernel=False, parent=app)
106+
yield KernelSpecManager(ensure_native_kernel=False)
117107

118108

119109
@pytest.fixture
120110
def kernel_spec_cache(is_enabled, kernel_spec_manager):
121-
122111
kspec_cache = KernelSpecCache.instance(
123-
kernel_spec_manager=kernel_spec_manager,
124-
cache_enabled=is_enabled,
125-
parent=kernel_spec_manager.parent,
112+
kernel_spec_manager=kernel_spec_manager, cache_enabled=is_enabled
126113
)
127114
yield kspec_cache
128-
kspec_cache.clear_instance()
115+
kspec_cache = None
116+
KernelSpecCache.clear_instance()
129117

130118

131119
@pytest.fixture(params=[False, True]) # Add types as needed
@@ -146,14 +134,12 @@ async def tests_get_named_spec(kernel_spec_cache):
146134
async def tests_get_modified_spec(kernel_spec_cache):
147135
kspec = await kernel_spec_cache.get_kernel_spec("test2")
148136
assert kspec.display_name == "Test kernel: test2"
149-
assert set(kspec.metadata['client_envs']) == set(UNION_CLIENT_ENVS)
150137

151138
# Modify entry
152139
_modify_kernelspec(kspec.resource_dir, "test2")
153140
await asyncio.sleep(0.5) # sleep for a half-second to allow cache to update item
154141
kspec = await kernel_spec_cache.get_kernel_spec("test2")
155142
assert kspec.display_name == "test2 modified!"
156-
assert set(kspec.metadata['client_envs']) == set(UNION_CLIENT_ENVS)
157143

158144

159145
async def tests_add_spec(kernel_spec_cache, kernelspec_location, other_kernelspec_location):

0 commit comments

Comments
 (0)