add refresh_user_hook

docs/source/how-to/refresh.md

@@ -132,3 +132,38 @@ c.Authenticator.auth_refresh_age = 0
 in which case the new `refresh_user` method will not be called.
 This is equivalent to the behavior of OAuthenticator 17.1 and earlier,
 where the default `refresh_user` was called, but did nothing.
+
+## Customizing refresh behavior
+
+There is also a `OAuthenticator.refresh_user_hook` configuration option,
+which allows you to override the refresh_user behavior.
+
+The hook is called as:
+
+```python
+refreshed = await refresh_user_hook(authentiator, user, auth_state)
+```
+
+where `refreshed` can be:
+
+- `True` if the user auth is up-to-date and nothing should change
+- `False` if the user should be forced to login again before they can do anything
+- `auth_data` - a dictionary containing the user model with that should be updated (see [`refresh_user`](inv:jupyterhub:py:method#jupyterhub.auth.Authenticator.refresh_user) docs)
+- `None` if the default `refresh_user` behavior should proceed
+
+For example, to use `refresh_user` for most users but have 'fake' users that don't exist in the oauth provider, you can return `True` for those users and None for others:
+
+```python
+infrastructure_users = {"health-check-user"}
+
+def refresh_user_hook(authenticator, user, auth_state):
+    if user.name in infrastructure_users:
+        # if this is an infrastructure user,
+        # refresh_user doesn't make sense
+        # consider it always fresh
+        return True
+    # for all other users, refresh as usual
+    return None
+
+c.OAuthenticator.refresh_user_hook = refresh_user_hook
+```

oauthenticator/oauth2.py

@@ -545,6 +545,32 @@ def _refresh_pre_spawn_default(self):
 
         return False
 
+    refresh_user_hook = Callable(
+        config=True,
+        default_value=None,
+        allow_none=True,
+        help="""
+        Hook for refreshing user auth info.
+
+        If given, allows overriding the `refresh_user` behavior.
+        Will be called as::
+
+            refreshed = await refresh_user_hook(authenticator, user, auth_state)
+
+        `refresh_user_hook` _may_ be async.
+
+        where `refreshed` can be:
+
+        - True (no change)
+        - False (require new login)
+        - auth_model (dict - the new auth model, if anything should be changeed)
+        - None (proceed with default refresh_user behavior -
+          allows overriding refresh_user behavior for _some_ users)
+
+        .. versionadded:: 17.3
+        """,
+    )
+
     logout_redirect_url = Unicode(
         config=True,
         help="""
@@ -1291,6 +1317,18 @@ async def authenticate(self, handler, data=None, **kwargs):
         # call the oauth endpoints
         return await self._token_to_auth_model(token_info)
 
+    async def _call_refresh_user_hook(self, user, auth_state):
+        """Call the refresh_user hook"""
+        try:
+            refreshed = self.refresh_user_hook(self, user, auth_state)
+            if isawaitable(refreshed):
+                refreshed = await refreshed
+        except Exception as e:
+            # let hook errors raise, nothing in auth should suppress errors
+            self.log.error(f"Error in refresh_user_hook: {e}")
+            raise
+        return refreshed
+
     async def refresh_user(self, user, handler=None, **kwargs):
         """
         Refresh user authentication
@@ -1325,7 +1363,14 @@ async def refresh_user(self, user, handler=None, **kwargs):
         if not self.enable_auth_state:
             # auth state not enabled, can't refresh
             return True
+
         auth_state = await user.get_auth_state()
+
+        if self.refresh_user_hook is not None:
+            refreshed = await self._call_refresh_user_hook(user, auth_state)
+            if refreshed is not None:
+                return refreshed
+
         if not auth_state:
             self.log.info(
                 f"No auth_state found for user {user.name} refresh, need full authentication",

oauthenticator/tests/test_generic.py

@@ -547,6 +547,22 @@ async def test_refresh_user(get_authenticator, generic_client, enable_refresh_to
     # from here on, enable auth state required for refresh to do anything
     authenticator.enable_auth_state = True
 
+    # case: custom refresh hook
+    async def async_hook(authenticator, user, auth_state):
+        return True
+
+    authenticator.refresh_user_hook = async_hook
+    refreshed = await authenticator.refresh_user(user, handler)
+    assert refreshed is True
+
+    def sync_hook(authenticator, user, auth_state):
+        return False
+
+    authenticator.refresh_user_hook = sync_hook
+    refreshed = await authenticator.refresh_user(user, handler)
+    assert refreshed is False
+    authenticator.refresh_user_hook = None
+
     # case: no auth state, but auth state enabled needs refresh
     auth_without_state = auth_model.copy()
     auth_without_state["auth_state"] = None