jupyterhub-cli.py

import urllib.request
import urllib.parse
import argparse
import json
import random
import time
import enum
import logging


logger = logging.getLogger(__name__)

DATA_PREFIX = b"data: "
RANDOM_REQUESTS_PER_MIN = 10


class State(enum.StrEnum):
    initial = enum.auto()
    check_status = enum.auto()
    start_server = enum.auto()
    wait_for_start = enum.auto()
    wait_for_stop = enum.auto()


def get_server_api_url(user_name: str, server_name: str | None) -> str:
    return (
        f"{api_url}/users/{user_name}/server"
        if server_name is None
        else f"{api_url}/users/{user_name}/servers/{server_name}"
    )


def run_loop(
    api_url: str,
    api_token: str,
    server_name: str | None,
    profile_options: dict[str, str],
):
    auth_headers = {"Authorization": f"token {api_token}"}

    # Initial state
    state = State.check_status

    # State data (side effects)
    user_name: str

    while True:
        logger.debug(state)
        match state:
            case State.check_status:
                # Get current user
                with urllib.request.urlopen(
                    urllib.request.Request(f"{api_url}/user", headers=auth_headers)
                ) as resp:
                    user_model = json.load(resp)

                # Side effect
                user_name = user_model["name"]

                # Is the server known of?
                existing_server = user_model["servers"].get(server_name or "")

                # Transition
                match existing_server:
                    case None:
                        state = State.start_server
                    case {"pending": "spawn", **rest}:
                        state = State.wait_for_start
                    case {"pending": "stop", **rest}:
                        state = State.wait_for_stop
                    case _:
                        assert existing_server["ready"]
                        return existing_server["url"]

            case State.start_server:
                # Try to start server
                with urllib.request.urlopen(
                    urllib.request.Request(
                        get_server_api_url(user_name, server_name),
                        method="POST",
                        data=json.dumps(profile_options).encode("utf-8"),
                        headers={**auth_headers, "Content-Type": "application/json"},
                    )
                ) as resp:
                    # Handle response
                    match resp.status:
                        # Server already running
                        case 201:
                            state = State.check_status
                        case 202:
                            state = State.wait_for_start
                        case 429:
                            retry_after = resp.headers.get("Retry-After")
                            delay = (
                                random.expovariate(RANDOM_REQUESTS_PER_MIN / 60)
                                if retry_after is None
                                else int(retry_after)
                            )
                            logger.info(
                                f"Server asked us to back off, waiting for {delay:.1f} seconds"
                            )
                            time.sleep(delay)

            case State.wait_for_start:
                # Get URL
                with urllib.request.urlopen(
                    urllib.request.Request(
                        f"{get_server_api_url(user_name, server_name)}/progress",
                        method="GET",
                        headers=auth_headers,
                    )
                ) as resp:
                    for line in iter(resp.readline, None):
                        if not line.startswith(DATA_PREFIX):
                            continue

                        logger.debug(line.decode())

                        # Load JSON response line
                        line_data = json.loads(line[len(DATA_PREFIX) :].decode())
                        if line_data.get("ready"):
                            return line_data["url"]

            case State.wait_for_stop:
                delay = random.expovariate(RANDOM_REQUESTS_PER_MIN / 60)
                logger.info(f"Waiting for server to stop for {delay:.1f} seconds")
                time.sleep(delay)
                state = State.check_status


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("url", help="JupyterHub API URL")
    parser.add_argument("token", help="JupyterHub API Token")
    parser.add_argument("--server", help="Server name")
    parser.add_argument(
        "-v", "--verbose", help="Turn on verbose debugging", action="store_true"
    )
    parser.add_argument(
        "--profile-option",
        help="Profile option key-value pair of the form key=value",
        nargs="*",
        default=[],
    )
    args = parser.parse_args()
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)

    # Arg processing
    server_name = args.server
    api_url = args.url.rstrip("/")
    api_token = args.token
    profile_options = {k: v for k, v in (p.split("=") for p in args.profile_option)}

    server_url = run_loop(api_url, api_token, server_name, profile_options)
    print(repr(server_url))