add heartbeat extension (jupyter-server#106)

* add heartbeat extension

* Bump to 0.8.1

Author: Zsailer

conftest.py

@@ -71,6 +71,7 @@ def jp_server_config(datastudio_env):
             "jpserver_extensions": {
                 "data_studio_jupyter_extensions": True,
                 "data_studio_jupyter_extensions.extensions.telemetry": True,
+                "data_studio_jupyter_extensions.extensions.heartbeat": True,
                 "nbclassic": True,
             },
         },
@@ -85,12 +86,3 @@ def datastudio_extension(jp_serverapp):
     point = pkg.extension_points["data_studio_jupyter_extensions"]
     app = point.app
     return app
-
-
-@pytest.fixture
-def datastudio_telemetry_extension(jp_serverapp):
-    name = "data_studio_jupyter_extensions.extensions.telemetry"
-    pkg = jp_serverapp.extension_manager.extensions[name]
-    point = pkg.extension_points["telemetry"]
-    app = point.app
-    return app

data_studio_jupyter_extensions/__init__.py

@@ -15,10 +15,24 @@ def _jupyter_labextension_paths():  # pragma: no cover
 
 def _jupyter_server_extension_points():  # pragma: no cover
     from data_studio_jupyter_extensions.app import DataStudioJupyterExtensions
+    from data_studio_jupyter_extensions.extensions.telemetry.extension import (
+        TelemetryExtension,
+    )
+    from data_studio_jupyter_extensions.extensions.heartbeat.extension import (
+        HeartbeatExtension,
+    )
 
     return [
         {
             "module": "data_studio_jupyter_extensions.app",
             "app": DataStudioJupyterExtensions,
-        }
+        },
+        {
+            "module": "data_studio_jupyter_extensions.extensions.telemetry.extension",
+            "app": TelemetryExtension,
+        },
+        {
+            "module": "data_studio_jupyter_extensions.extensions.heartbeat.extension",
+            "app": HeartbeatExtension,
+        },
     ]

data_studio_jupyter_extensions/_version.py

@@ -1 +1 @@
-__version__ = "0.8.0"  # pragma: no cover
+__version__ = "0.8.1"  # pragma: no cover

data_studio_jupyter_extensions/app.py

@@ -86,6 +86,11 @@ class DataStudioJupyterExtensions(ExtensionAppJinjaMixin, ExtensionApp):
     def _default_ssl_cert_file(self):  # pragma: no cover
         return get_ssl_cert(jupyter_runtime_dir())
 
+    dsw_url = UnicodeFromEnv(
+        name=constants.DSW_URL,
+        allow_none=True,
+    ).tag(config=True)
+
     datastudio_url = UnicodeFromEnv(
         name=constants.DS_URL, allow_none=True, help="ACS Data Platform URL."
     ).tag(config=True)
@@ -177,6 +182,7 @@ def _jupyter_server_config(self):
                     "ds_ws_max_buffer_size": self.ds_ws_max_buffer_size,
                     "ds_ws_max_message_size": self.ds_ws_max_message_size,
                 },
+                "allow_origin": self.dsw_url,
             },
         }
         apply_defaults(config, base_config)

data_studio_jupyter_extensions/constants.py

@@ -14,3 +14,4 @@
 STDINPORT = "STDINPORT"
 CONTROLPORT = "CONTROLPORT"
 NOTEBOOK_SERVICE_MODE = "NOTEBOOK_SERVICE_MODE"
+DSW_URL = "DSW_URL"

data_studio_jupyter_extensions/extensions/heartbeat/__init__.py

@@ -0,0 +1,12 @@
+from data_studio_jupyter_extensions.extensions.heartbeat.extension import (
+    HeartbeatExtension,
+)
+
+
+def _jupyter_server_extension_points():  # pragma: no cover
+    return [
+        {
+            "module": "data_studio_jupyter_extensions.extensions.heartbeat.extension",
+            "app": HeartbeatExtension,
+        }
+    ]

data_studio_jupyter_extensions/extensions/heartbeat/extension.py

@@ -0,0 +1,12 @@
+from jupyter_server.extension.application import ExtensionApp
+
+from .handlers import handlers
+
+
+class HeartbeatExtension(ExtensionApp):
+    """Jupyter Server extension that verifies
+    the health of the server.
+    """
+
+    name = "heartbeat"
+    handlers = handlers

data_studio_jupyter_extensions/extensions/heartbeat/handlers.py

@@ -0,0 +1,31 @@
+"""Application Heartbeat for PIE Health check"""
+import os
+from urllib.parse import urlparse
+
+from jupyter_server.base.handlers import JupyterHandler
+from jupyter_server.utils import url_path_join
+
+
+class HeartbeatHandler(JupyterHandler):
+    """Heartbeat handler"""
+
+    def set_default_headers(self):
+        datastudio_url = os.environ.get("DSW_URL")
+        parsed_uri = urlparse(datastudio_url)
+        domain = "{uri.scheme}://{uri.netloc}".format(uri=parsed_uri)
+        self.set_header("Access-Control-Allow-Origin", domain)
+        self.set_header("Access-Control-Allow-Headers", "Authorization")
+        self.set_header("Access-Control-Allow-Methods", "GET, OPTIONS")
+
+    def options(self):
+        self.set_status(204)
+        self.finish()
+
+    def get(self):
+        """HTTP GET Heartbeat handler"""
+        self.finish("Yep, we're healthy here.")
+
+
+handlers = [
+    (r"/__health", HeartbeatHandler),
+]

data_studio_jupyter_extensions/extensions/telemetry/__init__.py

@@ -4,10 +4,6 @@
 
 
 def _jupyter_server_extension_points():  # pragma: no cover
-    from data_studio_jupyter_extensions.extensions.telemetry.extension import (
-        TelemetryExtension,
-    )
-
     return [
         {
             "module": "data_studio_jupyter_extensions.extensions.telemetry.extension",

data_studio_jupyter_extensions/extensions/telemetry/extension.py

@@ -22,7 +22,11 @@ def initialize_settings(self):
         self.initialize_configurables()
         for schema_file in get_schema_files():
             self.telemetry_bus.register_schema_file(schema_file)
-        self.settings.update({"telemetry_bus": self.telemetry_bus})
+        self.settings.update(
+            {
+                "telemetry_bus": self.telemetry_bus,
+            }
+        )
 
     def initialize_configurables(self):
         handlers = []