From 91a290fda4c54a7008144d5b7e81e8d2b5716a45 Mon Sep 17 00:00:00 2001 From: Viacheslav Savinkov Date: Fri, 19 Jan 2024 19:56:43 -0500 Subject: [PATCH] Support named servers --- jupyterhub_ssh/__init__.py | 58 +++++++++++++++++++++++++++++--------- 1 file changed, 44 insertions(+), 14 deletions(-) diff --git a/jupyterhub_ssh/__init__.py b/jupyterhub_ssh/__init__.py index 6415113..4eb6389 100644 --- a/jupyterhub_ssh/__init__.py +++ b/jupyterhub_ssh/__init__.py @@ -34,29 +34,41 @@ def connection_made(self, conn): def password_auth_supported(self): return True - async def get_user_server_url(self, session, username): + async def get_user_server_url(self, session, username, force_default=False): """ Return user's server url if it is running. Else return None """ + server_name = "" + if "-" in username and not force_default: + username, server_name = username.split("-", 1) + print("Using username for custom instance", username, server_name) + else: + print("Using username for default instance", username) async with session.get(self.app.hub_url / "hub/api/users" / username) as resp: if resp.status != 200: return None user = await resp.json() - print(user) - # URLs will have preceding slash, but yarl forbids those server = user.get("servers", {}).get("", {}) if server.get("ready", False): - return self.app.hub_url / user["servers"][""]["url"][1:] + return self.app.hub_url / user["servers"][server_name]["url"][1:] else: return None - async def start_user_server(self, session, username): + async def start_user_server(self, session, username, force_default=False): """ """ # REST API reference: https://jupyterhub.readthedocs.io/en/stable/_static/rest-api/index.html#operation--users--name--server-post # REST API implementation: https://github.com/jupyterhub/jupyterhub/blob/187fe911edce06eb067f736eaf4cc9ea52e69e08/jupyterhub/apihandlers/users.py#L451-L497 create_url = self.app.hub_url / "hub/api/users" / username / "server" + server_name = "" + if "-" in username and not force_default: + username, server_name = username.split("-", 1) + + if server_name: + create_url = self.app.hub_url / "hub/api/users" / username / "servers" / server_name + + print("create_url: ", create_url, flush=True) async with session.post(create_url) as resp: if resp.status == 201 or resp.status == 400: @@ -68,6 +80,10 @@ async def start_user_server(self, session, username): # We manually generate this, even though it's *bad* # Mostly because when the server is already running, JupyterHub # doesn't respond with the whole model! + if server_name: + print("Custom-named server started", username, server_name) + return self.app.hub_url / "user" / username / server_name + print("Server started", username, server_name) return self.app.hub_url / "user" / username elif resp.status == 202: # Server start has been requested, now and potentially earlier, @@ -78,35 +94,49 @@ async def start_user_server(self, session, username): async with timeout(self.app.start_timeout): notebook_url = None self._conn.send_auth_banner("Starting your server...") + print("Server start requested", username, server_name) while notebook_url is None: # FIXME: Exponential backoff + make this configurable await asyncio.sleep(0.5) - notebook_url = await self.get_user_server_url( - session, username - ) + if server_name: + notebook_url = await self.get_user_server_url( + session, username + "/" + server_name, force_default + ) + else: + notebook_url = await self.get_user_server_url( + session, username, force_default + ) self._conn.send_auth_banner(".") self._conn.send_auth_banner("done!\n") + print("Server started", username, server_name) return notebook_url except asyncio.TimeoutError: # Server didn't start on time! self._conn.send_auth_banner("failed to start server on time!\n") + print("Server start timed out") return None elif resp.status == 403: - # Token is wrong! + print(f"The token seems to be wrong for {username}") return None + elif resp.status == 404 and server_name: + print(f"Custom instance not found, probably {username} contains dash") + return 404 else: # FIXME: Handle other cases that pop up + print("Unhandled response", resp.status) resp.raise_for_status() - async def validate_password(self, username, token): + async def validate_password(self, username, token, force_default=False): self.username = username self.token = token headers = {"Authorization": f"token {token}"} async with ClientSession(headers=headers) as session: - notebook_url = await self.start_user_server(session, username) + notebook_url = await self.start_user_server(session, username, force_default) if notebook_url is None: return False + elif notebook_url == 404: + return self.validate_password(username, token, force_default=True) else: self.notebook_url = notebook_url return True @@ -225,13 +255,13 @@ class JupyterHubSSH(Application): debug = Bool( True, help=""" - Turn on debugg logging + Turn on debug logging """, config=True, ) hub_url = Any( - "", + URL("http://proxy-public.jhub.svc.cluster.local"), help=""" URL of JupyterHub's proxy to connect to. @@ -269,7 +299,7 @@ def _hub_url_cast_string_to_yarl_url(self, proposal): raise ValueError("hub_url must either be a string or a yarl.URL") host_key_path = Unicode( - "", + "/etc/jupyterhub-ssh/config/hostKey", help=""" Path to host's private SSH Key.