Add support of using custom env variables

jupyter_server/serverapp.py

@@ -401,6 +401,8 @@ def init_settings(
             # collapse $HOME to ~
             root_dir = "~" + root_dir[len(home) :]
 
+        self.allow_custom_env_variables = jupyter_app.allow_custom_env_variables
+
         settings = {
             # basics
             "log_function": log_request,
@@ -460,6 +462,7 @@ def init_settings(
             "server_root_dir": root_dir,
             "jinja2_env": env,
             "serverapp": jupyter_app,
+            "page_config_hook": (self.page_config_hook),
         }
 
         # allow custom overrides for the tornado web app.
@@ -470,6 +473,10 @@ def init_settings(
             settings["xsrf_cookie_kwargs"] = {"path": base_url}
         return settings
 
+    def page_config_hook(self, handler, page_config):
+        page_config["allow_custom_env_variables"] = self.allow_custom_env_variables
+        return page_config
+
     def init_handlers(self, default_services, settings):
         """Load the (URL pattern, handler) tuples for each component."""
         # Order matters. The first handler to match the URL will handle the request.
@@ -1427,6 +1434,12 @@ def _default_allow_remote(self) -> bool:
                         """,
     )
 
+    allow_custom_env_variables = Bool(
+        False,
+        config=True,
+        help="""Allow to use custom env variables""",
+    )
+
     browser = Unicode(
         "",
         config=True,

jupyter_server/services/kernels/kernelmanager.py

@@ -224,6 +224,7 @@ async def _async_start_kernel(  # type:ignore[override]
             The name identifying which kernel spec to launch. This is ignored if
             an existing kernel is returned, but it may be checked in the future.
         """
+
         if kernel_id is None or kernel_id not in self:
             if path is not None:
                 kwargs["cwd"] = self.cwd_for_path(path, env=kwargs.get("env", {}))

jupyter_server/services/sessions/handlers.py

@@ -77,6 +77,7 @@ async def post(self):
         kernel = model.get("kernel", {})
         kernel_name = kernel.get("name", None)
         kernel_id = kernel.get("id", None)
+        custom_env_vars = kernel.get("custom_env_vars", {})
 
         if not kernel_id and not kernel_name:
             self.log.debug("No kernel specified, using default kernel")
@@ -93,6 +94,7 @@ async def post(self):
                     kernel_id=kernel_id,
                     name=name,
                     type=mtype,
+                    custom_env_vars=custom_env_vars,
                 )
             except NoSuchKernel:
                 msg = (
@@ -152,6 +154,8 @@ async def patch(self, session_id):
             changes["name"] = model["name"]
         if "type" in model:
             changes["type"] = model["type"]
+        if "custom_env_vars" in model:
+            changes["custom_env_vars"] = model["custom_env_vars"]
         if "kernel" in model:
             # Kernel id takes precedence over name.
             if model["kernel"].get("id") is not None:
@@ -160,13 +164,19 @@ async def patch(self, session_id):
                     raise web.HTTPError(400, "No such kernel: %s" % kernel_id)
                 changes["kernel_id"] = kernel_id
             elif model["kernel"].get("name") is not None:
+                if "custom_env_vars" in model["kernel"]:
+                    custom_env_vars = model["kernel"]["custom_env_vars"]
+                else:
+                    custom_env_vars = {}
+
                 kernel_name = model["kernel"]["name"]
                 kernel_id = await sm.start_kernel_for_session(
                     session_id,
                     kernel_name=kernel_name,
                     name=before["name"],
                     path=before["path"],
                     type=before["type"],
+                    custom_env_vars=custom_env_vars,
                 )
                 changes["kernel_id"] = kernel_id
 

jupyter_server/services/sessions/sessionmanager.py

@@ -209,6 +209,7 @@ def __init__(self, *args, **kwargs):
     _cursor = None
     _connection = None
     _columns = {"session_id", "path", "name", "type", "kernel_id"}
+    _custom_envs = {}
 
     @property
     def cursor(self):
@@ -267,6 +268,7 @@ async def create_session(
         type: Optional[str] = None,
         kernel_name: Optional[KernelName] = None,
         kernel_id: Optional[str] = None,
+        custom_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """Creates a session and returns its model
 
@@ -283,7 +285,7 @@ async def create_session(
             pass
         else:
             kernel_id = await self.start_kernel_for_session(
-                session_id, path, name, type, kernel_name
+                session_id, path, name, type, kernel_name, custom_env_vars
             )
         record.kernel_id = kernel_id
         self._pending_sessions.update(record)
@@ -319,6 +321,7 @@ async def start_kernel_for_session(
         name: Optional[ModelName],
         type: Optional[str],
         kernel_name: Optional[KernelName],
+        custom_env_vars: Optional[Dict[str, Any]] = None,
     ) -> str:
         """Start a new kernel for a given session.
 
@@ -335,16 +338,25 @@ async def start_kernel_for_session(
             the type of the session
         kernel_name : str
             the name of the kernel specification to use.  The default kernel name will be used if not provided.
+        custom_env_vars: dict
+            dictionary of custom env variables
         """
         # allow contents manager to specify kernels cwd
         kernel_path = await ensure_async(self.contents_manager.get_kernel_path(path=path))
-
         kernel_env = self.get_kernel_env(path, name)
+
+        # if we have custom env than we have to add them to available env variables
+        if custom_env_vars is not None and isinstance(custom_env_vars, dict):
+            for key, value in custom_env_vars.items():
+                kernel_env[key] = value
+
         kernel_id = await self.kernel_manager.start_kernel(
             path=kernel_path,
             kernel_name=kernel_name,
             env=kernel_env,
         )
+
+        self._custom_envs[kernel_id] = custom_env_vars
         return cast(str, kernel_id)
 
     async def save_session(self, session_id, path=None, name=None, type=None, kernel_id=None):
@@ -445,6 +457,7 @@ async def update_session(self, session_id, **kwargs):
             and the value replaces the current value in the session
             with session_id.
         """
+
         await self.get_session(session_id=session_id)
 
         if not kwargs:
@@ -464,7 +477,18 @@ async def update_session(self, session_id, **kwargs):
                 "SELECT path, name, kernel_id FROM session WHERE session_id=?", [session_id]
             )
             path, name, kernel_id = self.cursor.fetchone()
-            self.kernel_manager.update_env(kernel_id=kernel_id, env=self.get_kernel_env(path, name))
+
+            env = self.get_kernel_env(path, name)
+
+            # if we have custom env than we have to add them to available env variables
+            if self._custom_envs is not None and isinstance(self._custom_envs, dict):
+                if self._custom_envs[kernel_id] is not None and isinstance(
+                    self._custom_envs[kernel_id], dict
+                ):
+                    for key, value in self._custom_envs[kernel_id].items():
+                        env[key] = value
+
+            self.kernel_manager.update_env(kernel_id=kernel_id, env=env)
 
     async def kernel_culled(self, kernel_id: str) -> bool:
         """Checks if the kernel is still considered alive and returns true if its not found."""
@@ -526,6 +550,9 @@ async def delete_session(self, session_id):
         record = KernelSessionRecord(session_id=session_id)
         self._pending_sessions.update(record)
         session = await self.get_session(session_id=session_id)
-        await ensure_async(self.kernel_manager.shutdown_kernel(session["kernel"]["id"]))
+        kernel_id = session["kernel"]["id"]
+        if kernel_id in self._custom_envs:
+            del self._custom_envs[kernel_id]
+        await ensure_async(self.kernel_manager.shutdown_kernel(kernel_id))
         self.cursor.execute("DELETE FROM session WHERE session_id=?", (session_id,))
         self._pending_sessions.remove(record)