diff --git a/jupyterhub/apihandlers/users.py b/jupyterhub/apihandlers/users.py index 01356eb0..eb35a13f 100644 --- a/jupyterhub/apihandlers/users.py +++ b/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 {0}'.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 {0}'.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? diff --git a/jupyterhub/spawner.py b/jupyterhub/spawner.py index 584f5c45..99faf081 100644 --- a/jupyterhub/spawner.py +++ b/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 diff --git a/jupyterhub/tests/test_api.py b/jupyterhub/tests/test_api.py index 7a079e45..b7002e79 100644 --- a/jupyterhub/tests/test_api.py +++ b/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 {0}'.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