Merge pull request #4455 from kreuzert/main

Add Spawner.progress_ready_hook for customizing the ready event

jupyterhub/apihandlers/users.py

@@ -2,6 +2,7 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 import asyncio
+import inspect
 import json
 from datetime import datetime, timedelta, timezone
 
@@ -710,19 +711,32 @@ class SpawnProgressAPIHandler(APIHandler):
         # - spawner not running at all
         # - spawner failed
         # - spawner pending start (what we expect)
-        url = url_path_join(user.url, url_escape_path(server_name), '/')
-        ready_event = {
-            'progress': 100,
-            'ready': True,
-            'message': f"Server ready at {url}",
-            'html_message': 'Server ready at <a href="{0}">{0}</a>'.format(url),
-            'url': url,
-        }
         failed_event = {'progress': 100, 'failed': True, 'message': "Spawn failed"}
 
+        async def get_ready_event():
+            url = url_path_join(user.url, url_escape_path(server_name), '/')
+            ready_event = {
+                'progress': 100,
+                'ready': True,
+                'message': f"Server ready at {url}",
+                'html_message': 'Server ready at <a href="{0}">{0}</a>'.format(url),
+                'url': url,
+            }
+            original_ready_event = ready_event.copy()
+            if spawner.progress_ready_hook:
+                try:
+                    ready_event = spawner.progress_ready_hook(spawner, ready_event)
+                    if inspect.isawaitable(ready_event):
+                        ready_event = await ready_event
+                except Exception as e:
+                    self.log.exception(f"Error in ready_event hook: {e}")
+                    ready_event = original_ready_event
+            return ready_event
+
         if spawner.ready:
             # spawner already ready. Trigger progress-completion immediately
             self.log.info("Server %s is already started", spawner._log_name)
+            ready_event = await get_ready_event()
             await self.send_event(ready_event)
             return
 
@@ -766,6 +780,7 @@ class SpawnProgressAPIHandler(APIHandler):
         if spawner.ready:
             # spawner is ready, signal completion and redirect
             self.log.info("Server %s is ready", spawner._log_name)
+            ready_event = await get_ready_event()
             await self.send_event(ready_event)
         else:
             # what happened? Maybe spawn failed?

jupyterhub/spawner.py

@@ -840,6 +840,27 @@ class Spawner(LoggingConfigurable):
         """,
     ).tag(config=True)
 
+    progress_ready_hook = Any(
+        help="""
+        An optional hook function that you can implement to modify the
+        ready event, which will be shown to the user on the spawn progress page when their server
+        is ready.
+
+        This can be set independent of any concrete spawner implementation.
+
+        This maybe a coroutine.
+
+        Example::
+
+            async def my_ready_hook(spawner, ready_event):
+                ready_event["html_message"] = f"Server {spawner.name} is ready for {spawner.user.name}"
+                return ready_event
+
+            c.Spawner.progress_ready_hook = my_ready_hook
+
+        """
+    ).tag(config=True)
+
     pre_spawn_hook = Any(
         help="""
         An optional hook function that you can implement to do some

jupyterhub/tests/test_api.py

@@ -1145,6 +1145,92 @@ async def test_progress_ready(request, app):
     assert evt['url'] == app_user.url
 
 
+async def test_progress_ready_hook_async_func(request, app):
+    """Test progress ready hook in Spawner class with an async function"""
+    db = app.db
+    name = 'saga'
+    app_user = add_user(db, app=app, name=name)
+    html_message = 'customized html message'
+    spawner = app_user.spawner
+
+    async def custom_progress_ready_hook(spawner, ready_event):
+        ready_event['html_message'] = html_message
+        return ready_event
+
+    spawner.progress_ready_hook = custom_progress_ready_hook
+    r = await api_request(app, 'users', name, 'server', method='post')
+    r.raise_for_status()
+    r = await api_request(app, 'users', name, 'server/progress', stream=True)
+    r.raise_for_status()
+    request.addfinalizer(r.close)
+    assert r.headers['content-type'] == 'text/event-stream'
+    ex = async_requests.executor
+    line_iter = iter(r.iter_lines(decode_unicode=True))
+    evt = await ex.submit(next_event, line_iter)
+    assert evt['progress'] == 100
+    assert evt['ready']
+    assert evt['url'] == app_user.url
+    assert evt['html_message'] == html_message
+
+
+async def test_progress_ready_hook_sync_func(request, app):
+    """Test progress ready hook in Spawner class with a sync function"""
+    db = app.db
+    name = 'saga'
+    app_user = add_user(db, app=app, name=name)
+    html_message = 'customized html message'
+    spawner = app_user.spawner
+
+    def custom_progress_ready_hook(spawner, ready_event):
+        ready_event['html_message'] = html_message
+        return ready_event
+
+    spawner.progress_ready_hook = custom_progress_ready_hook
+    r = await api_request(app, 'users', name, 'server', method='post')
+    r.raise_for_status()
+    r = await api_request(app, 'users', name, 'server/progress', stream=True)
+    r.raise_for_status()
+    request.addfinalizer(r.close)
+    assert r.headers['content-type'] == 'text/event-stream'
+    ex = async_requests.executor
+    line_iter = iter(r.iter_lines(decode_unicode=True))
+    evt = await ex.submit(next_event, line_iter)
+    assert evt['progress'] == 100
+    assert evt['ready']
+    assert evt['url'] == app_user.url
+    assert evt['html_message'] == html_message
+
+
+async def test_progress_ready_hook_async_func_exception(request, app):
+    """Test progress ready hook in Spawner class with an exception in
+    an async function
+    """
+    db = app.db
+    name = 'saga'
+    app_user = add_user(db, app=app, name=name)
+    html_message = 'Server ready at <a href="{0}">{0}</a>'.format(app_user.url)
+    spawner = app_user.spawner
+
+    async def custom_progress_ready_hook(spawner, ready_event):
+        ready_event["html_message"] = "."
+        raise Exception()
+
+    spawner.progress_ready_hook = custom_progress_ready_hook
+    r = await api_request(app, 'users', name, 'server', method='post')
+    r.raise_for_status()
+    r = await api_request(app, 'users', name, 'server/progress', stream=True)
+    r.raise_for_status()
+    request.addfinalizer(r.close)
+    assert r.headers['content-type'] == 'text/event-stream'
+    ex = async_requests.executor
+    line_iter = iter(r.iter_lines(decode_unicode=True))
+    evt = await ex.submit(next_event, line_iter)
+    assert evt['progress'] == 100
+    assert evt['ready']
+    assert evt['url'] == app_user.url
+    assert evt['html_message'] == html_message
+
+
 async def test_progress_bad(request, app, bad_spawn):
     """Test progress API when spawner has already failed"""
     db = app.db