Log events that link kernels to client sessions (jupyter-server#91)

* add custom session manager and begin logging events

* add handler for debug mode that prints to standout

* add a test for telemetry

* add event before session starts

* add nbclassic as an explicit extension

* uncomment tests

conftest.py

@@ -3,6 +3,7 @@
 from jupyter_server.serverapp import ServerApp
 
 from data_studio_jupyter_extensions import constants
+from data_studio_jupyter_extensions.app import DataStudioJupyterExtensions
 from data_studio_jupyter_extensions.notebook_service import NotebookServiceClient
 from data_studio_jupyter_extensions.telemetry.logger import TelemetryBus
 from data_studio_jupyter_extensions.tests.mock.utils import load_openapi_spec
@@ -62,10 +63,13 @@ def clear_singletons():
 
 
 @pytest.fixture
-def jp_server_config():
+def jp_server_config(datastudio_env):
     config = {
         "ServerApp": {
-            "jpserver_extensions": {"data_studio_jupyter_extensions": True},
+            "jpserver_extensions": {
+                "data_studio_jupyter_extensions": True,
+                "nbclassic": True,
+            },
         },
     }
     return config

data_studio_jupyter_extensions/app.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import pathlib
 import sys
@@ -162,7 +163,9 @@ def _jupyter_server_config(self):
             },
             "Session": {"key": b"notebooks"},
             "ServerApp": {
+                "log_level": "DEBUG",
                 "kernel_spec_manager_class": "data_studio_jupyter_extensions.kernelspecs.DSKernelSpecManager",
+                "session_manager_class": "data_studio_jupyter_extensions.session_manager.DSSessionManager",
                 "login_handler_class": "data_studio_jupyter_extensions.auth.login.JWTLoginHandler",
                 "logout_handler_class": "data_studio_jupyter_extensions.auth.logout.DSLogoutHandler",
                 "cookie_options": {"expires_days": 1},
@@ -218,8 +221,17 @@ def initialize_configurables(self):
             project_id=self.project_id,
             notebook_id=self.notebook_id,
         )
+        handlers = []
+        if self.serverapp.log_level <= 10:
+            handler = logging.StreamHandler(stream=sys.stdout)
+            handlers.append(handler)
         self.telemetry_bus = TelemetryBus.instance(
-            parent=self, allowed_schemas=["event.datastudio.jupyter.com/kernel-message"]
+            parent=self,
+            handlers=handlers,
+            allowed_schemas=[
+                "event.datastudio.jupyter.com/kernel-message",
+                "event.datastudio.jupyter.com/session-message",
+            ],
         )
 
 

data_studio_jupyter_extensions/session_manager.py

@@ -0,0 +1,54 @@
+import os
+import uuid
+
+from jupyter_server.services.sessions.sessionmanager import SessionManager
+from traitlets import default
+from traitlets import Instance
+
+from data_studio_jupyter_extensions.telemetry.logger import TelemetryBus
+
+
+class DSSessionManager(SessionManager):
+
+    telemetry_bus = Instance(TelemetryBus)
+
+    @default("telemetry_bus")
+    def _default_telemetry_bus(self):  # pragma: no cover
+        return TelemetryBus.instance(parent=self.parent)
+
+    async def start_kernel_for_session(self, session_id, path, name, type, kernel_name):
+        """
+        Start a new kernel for a given session.
+
+        This overrides the base Session Manager
+        """
+        # Get a kernel ID earliers than
+        kernel_id = str(uuid.uuid4())
+        # allow contents manager to specify kernels cwd
+        kernel_path = self.contents_manager.get_kernel_path(path=path)
+        # Emit the Session information.
+        message = f"Starting a session for kernel: {kernel_id}."
+        self.telemetry_bus.record_event(
+            schema_name="event.datastudio.jupyter.com/session-message",
+            version=1,
+            event={
+                "kernel_path": os.path.join(kernel_path, name),
+                "kernel_id": kernel_id,
+                "message": message,
+            },
+        )
+        kernel_id = await self.kernel_manager.start_kernel(
+            kernel_id=kernel_id, path=kernel_path, kernel_name=kernel_name
+        )
+        # Emit the Session information.
+        message = "Session created for kernel."
+        self.telemetry_bus.record_event(
+            schema_name="event.datastudio.jupyter.com/session-message",
+            version=1,
+            event={
+                "kernel_path": os.path.join(kernel_path, name),
+                "kernel_id": kernel_id,
+                "message": message,
+            },
+        )
+        return kernel_id

data_studio_jupyter_extensions/telemetry/schemas/sessions/session-message.v1.yaml

@@ -0,0 +1,24 @@
+$id: event.datastudio.jupyter.com/session-message
+version: 1
+title: Session Message
+description: |
+  Emit a message from the kernel provisioner
+type: object
+properties:
+  kernel_path:
+    title: Kernel Session Path
+    description: |
+      Path that links a running kernel to a given frontend
+      session (e.g. notebook, console, etc.).
+  kernel_id:
+    title: Kernel ID
+    description: |
+      The UUID for a given kernel.
+  message:
+    title: Message type
+    description: |
+      Message returned by the Session manager.
+required:
+  - kernel_path
+  - kernel_id
+  - message

data_studio_jupyter_extensions/tests/test_telemetry.py

@@ -0,0 +1,31 @@
+import io
+import json
+import logging
+
+import pytest
+
+from data_studio_jupyter_extensions.telemetry.logger import TelemetryBus
+
+
+async def test_telemetry_websocket(datastudio_extension, jp_ws_fetch):
+    # Get telemetry_bus from extension
+    telemetry_bus = datastudio_extension.telemetry_bus
+
+    # Open a websocket connection.
+    ws = await jp_ws_fetch("/subscribe")
+
+    telemetry_bus.record_event(
+        schema_name="event.datastudio.jupyter.com/kernel-message",
+        version=1,
+        event={"process_id": "test process id", "message": "test"},
+    )
+
+    message = await ws.read_message()
+    event_data = json.loads(message)
+    # Close websocket
+    ws.close()
+
+    assert "process_id" in event_data
+    assert event_data["process_id"] == "test process id"
+    assert "message" in event_data
+    assert event_data["message"] == "test"