How to access current user/auth_state from notebook

Question

I am using JupyterHub with custom authenticator. It sets auth_state with access token, which then can be copied into environment inside pre_spawn_start method like in the example:

class MyAuthenticator(Authenticator):
    @gen.coroutine
    def authenticate(self, handler, data=None):
        username = yield identify_user(handler, data)
        upstream_token = yield token_for_user(username)
        return {
            'name': username,
            'auth_state': {
                'upstream_token': upstream_token,
            },
        }

@gen.coroutine
def pre_spawn_start(self, user, spawner):
    """Pass upstream_token to spawner via environment variable"""
    auth_state = yield user.get_auth_state()
    if not auth_state:
        # auth_state not enabled
        return
    spawner.environment['UPSTREAM_TOKEN'] = auth_state['upstream_token']

However, it happens only once per user. If I logout and login again, the pre_spawn_start is not called again and old token is still present in the environment variable.

Is it possible to access user.get_auth_state() directly from the notebook so I can be sure to use the current token, rather then previously set and stored in the environment?

Otherwise, is it possible to force spawner stop on logout so the subsequent login will trigger the pre_spawn_start?

Answer 1

you can force stop servers and ask for login when token gets expired, it can achieved by following steps,

1. Write refresh_user method to get new token(attached sample code)
2. Set time for refresh_user.
   • c.MyAuthenticator.auth_refresh_age = 30 #This could be the time of validity of token in seconds

import inspect, concurrent,asyncio

async def refresh_user(self, user,handler=None):
    """
    1. Check if token is valid and then call _shutdown_servers and then redirect to login page
    2. If time of refresh_user is set as token expiry, directly call _shutdown_servers and then redirect to login page
    This is shutdown single user servers and once redirected to login, auth flow gets run and new tokens are passed to spawner
    """
    auth_state = await user.get_auth_state()
    if self._is_invalid_sessionid(auth_state):
        await self._shutdown_servers(user, handler=None)
        handler.clear_login_cookie()
        handler.redirect('/login')
    return True

async def maybe_future(obj):
    """Return an asyncio Future
    Use instead of gen.maybe_future
    For our compatibility, this must accept:
    - asyncio coroutine (gen.maybe_future doesn't work in tornado < 5)
    - tornado coroutine (asyncio.ensure_future doesn't work)
    - scalar (asyncio.ensure_future doesn't work)
    - concurrent.futures.Future (asyncio.ensure_future doesn't work)
    - tornado Future (works both ways)
    - asyncio Future (works both ways)
    """
    if inspect.isawaitable(obj):
        # already awaitable, use ensure_future
        return asyncio.ensure_future(obj)
    elif isinstance(obj, concurrent.futures.Future):
        return asyncio.wrap_future(obj)
    else:
        # could also check for tornado.concurrent.Future
        # but with tornado >= 5.1 tornado.Future is asyncio.Future
        f = asyncio.Future()
        f.set_result(obj)
        return f

async def _shutdown_servers(self, user,handler):
    """Shutdown servers for logout
    Get all active servers for the provided user, stop them.
    """
    active_servers = [
        name
        for (name, spawner) in user.spawners.items()
        if spawner.active and not spawner.pending
    ]
    if active_servers:
        self.log.info("Shutting down %s's servers", user.name)
        futures = []
        for server_name in active_servers:
            futures.append(self.maybe_future(handler.stop_single_user(user, server_name)))
        await asyncio.gather(*futures)

https://jupyterhub.readthedocs.io/en/stable/api/auth.html#jupyterhub.auth.Authenticator.auth_refresh_age https://jupyterhub.readthedocs.io/en/stable/api/auth.html#jupyterhub.auth.Authenticator.refresh_user

Answer 2

Seems this c.Authenticator.refresh_pre_spawn = True, as suggestes by Eos Antigen solves the question: everytime you log in again the pre-spawn defined function executes again.