Add Kernel restart, and move apple cert into Jupyter runtime dir (jupyter-server#57)

* enable kernel restart

* add some comments, cleanup, etc

* remove old env vars

Author: Zsailer

data_studio_jupyter_extensions/app.py

@@ -1,6 +1,7 @@
 import os
 import pathlib
 
+from jupyter_core.paths import jupyter_runtime_dir
 from jupyter_server.config_manager import recursive_update
 from jupyter_server.extension.application import ExtensionApp
 from jupyter_server.extension.application import ExtensionAppJinjaMixin
@@ -18,12 +19,23 @@
 from .telemetry.utils import DEFAULT_TEMPLATE_FILES_PATH
 from .telemetry.utils import get_schema_files
 from .traits import UnicodeFromEnv
+from .utils import get_ssl_cert
 
 
 MODE_CONFIG_PATH = pathlib.Path(__file__).parent.joinpath("config")
 
 
 class DataStudioJupyterExtensions(ExtensionAppJinjaMixin, ExtensionApp):
+    """JupyterLab with Data Studio extensions.
+
+    This extension bundles multiple Jupyter server and client extensions to
+    ship a JupyterLab experience that serves Apple Data Scientists.
+
+    The application can be launched by calling:
+    ```
+    DataStudioJupyterExtensions.launch_instance()
+    ```
+    """
 
     name = "data_studio_jupyter_extensions"
     default_url = "/lab"
@@ -32,18 +44,36 @@ class DataStudioJupyterExtensions(ExtensionAppJinjaMixin, ExtensionApp):
     template_paths = [str(DEFAULT_TEMPLATE_FILES_PATH)]
     static_paths = [str(DEFAULT_STATIC_FILES_PATH)]
 
+    # Add configurables to the list of classes displayed
+    # in the --help.
+    classes = [NotebookServiceClient, TelemetryBus]
+
     handlers = handlers
 
     aliases = {"mode": "DataStudioJupyterExtensions.mode"}
 
     mode = UnicodeFromEnv(
         default_value="local-local",
-        help="The type",
+        help="The mode describing how/where the application is being launched."
+        "There are three option, 'local-local', 'local-cluster', or "
+        "'cluster-cluster'. ",
         name=constants.NOTEBOOK_SERVICE_MODE,
     ).tag(config=True)
 
     mode_config_path = str(MODE_CONFIG_PATH)
 
+    ssl_cert_file = UnicodeFromEnv(
+        name="DS_SSL_CERT_FILE",
+        help="The Apple corporate SSL cert file necessary for accessing "
+        "the Data Platform UI. If not given, this application will try to "
+        "automatically download a new SSL cert to the Runtime directory.",
+        allow_none=True,
+    ).tag(config=True)
+
+    @default("ssl_cert_file")
+    def _default_ssl_cert_file(self):  # pragma: no cover
+        return get_ssl_cert(jupyter_runtime_dir())
+
     # List configurables here.
     nbservice_client_class = Type(
         default_value=NotebookServiceClient,
@@ -96,8 +126,7 @@ def initialize_settings(self):
 
     def initialize_configurables(self):
         self.nbservice_client = self.nbservice_client_class.instance(
-            parent=self,
-            log=self.log,
+            parent=self, log=self.log, ssl_cert_file=self.ssl_cert_file
         )
         self.telemetry_bus = TelemetryBus.instance(
             allowed_schemas=["event.datastudio.jupyter.com/kernel-message"]

data_studio_jupyter_extensions/notebook_service.py

@@ -3,6 +3,7 @@
 import ssl
 import urllib
 
+from jupyter_core.paths import jupyter_runtime_dir
 from jupyter_server.utils import url_path_join as ujoin
 from tornado.escape import json_decode
 from tornado.httpclient import AsyncHTTPClient
@@ -44,7 +45,7 @@ class NotebookServiceClient(SingletonConfigurable):
 
     @default("ssl_cert_file")
     def _default_ssl_cert_file(self):  # pragma: no cover
-        return get_ssl_cert()
+        return get_ssl_cert(jupyter_runtime_dir())
 
     http_client = Instance(AsyncHTTPClient).tag(config=True)
 

data_studio_jupyter_extensions/provisioner.py

@@ -337,6 +337,8 @@ async def terminate(self, restart: bool = False) -> None:
 
     async def shutdown_requested(self, restart: bool = False) -> None:
         await self.nbservice_client.stop_kernel(self.process_id)
+        if restart:
+            self.process_id = None
 
     async def pre_launch(self, **kwargs: Any) -> Dict[str, Any]:
         return await super().pre_launch(cmd=[], **kwargs)

data_studio_jupyter_extensions/traits.py

@@ -17,11 +17,21 @@ class FromEnvMixin:
     # and prevents the dynamic default from working.
     default_value = Undefined
 
-    def __init__(self, name, *args, **kwargs):
+    def __init__(self, name, help=None, *args, **kwargs):
         self.envvar_name = name
+        # Prefix the "help" string in these type of traits
+        # with some documentation that mentions the
+        # environment variable.
+        help_prefix = (
+            f"(This trait is also configurable from the environment variable, {name}.) "
+        )
+        if not help:
+            help = help_prefix
+        else:
+            help = help_prefix + help
         # The actual traittype is second in the MRO.
         traittype = self.__class__.mro()[2]
-        traittype.__init__(self, *args, **kwargs)
+        traittype.__init__(self, help=help, *args, **kwargs)
 
     def make_dynamic_default(self):
         # Make sure to dynamically load the default when the

data_studio_jupyter_extensions/utils.py

@@ -13,12 +13,13 @@
 from tenacity import stop_after_delay
 from tenacity import wait_exponential
 
+
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 
-def get_ssl_cert():
-    """Get Apple corproate SSL certificate, write
-    to a temporary file and return that file name.
+def get_ssl_cert(path):  # pragma: no cover
+    """Get Apple corporate SSL certificate, write
+    to `path`.
     """
     urls = [
         "https://certificatemanager.apple.com/certs/apple_corporate_root_ca.pem",
@@ -29,7 +30,7 @@ def get_ssl_cert():
         r = requests.get(url, verify=False)
         certs += r.text
 
-    fpath = pathlib.Path("apple_root_ca.pem")
+    fpath = pathlib.Path(path).joinpath("apple_root_ca.pem")
     fpath.write_text(certs)
     return str(fpath.absolute())
 

env.sh

@@ -6,6 +6,3 @@ export ISSUER=https://iam.corp.apple.com
 export KEY_URL=https://iam.corp.apple.com/oauth2/v3/certs
 export AUDIENCE=notebook-server-int
 export CLIENT_ID=d314ef93-a7b3-4451-b68d-51bd395ba7f2
-# export DS_PROCESS_ID=4p27l98hlgm9
-export DS_SPEC_ID=bkvgsydsi1rd
-export KG_LIST_KERNELS=True