Backport PR #3437: patch base handlers from both jupyter_server and notebook

and clarify warning when a base handler isn't patched that auth is still being applied

- reorganize patch steps into functions for easier re-use
- patch notebook and jupyter_server handlers if they are already imported
- run patch after initialize to ensure extensions have done their importing before we check what's present
- apply class-level patch even when instance-level patch is happening to avoid triggering patch on every request

This change isn't as big as it looks, because it's mostly moving some re-used code to a couple of functions.

closes https://github.com/jupyter-server/jupyter_server/issues/488

Signed-off-by: Min RK <benjaminrk@gmail.com>

jupyterhub/singleuser/mixins.py

@@ -10,9 +10,11 @@ with JupyterHub authentication mixins enabled.
 # Distributed under the terms of the Modified BSD License.
 import asyncio
 import json
+import logging
 import os
 import random
 import secrets
+import sys
 import warnings
 from datetime import datetime
 from datetime import timezone
@@ -99,19 +101,26 @@ class JupyterHubLoginHandlerMixin:
         Thus shouldn't be called anymore because HubAuthenticatedHandler
         should have already overridden get_current_user().
 
-        Keep here to prevent unlikely circumstance from losing auth.
+        Keep here to protect uncommon circumstance of multiple BaseHandlers
+        from missing auth.
+
+        e.g. when multiple BaseHandler classes are used.
         """
         if HubAuthenticatedHandler not in handler.__class__.mro():
             warnings.warn(
-                f"Expected to see HubAuthenticatedHandler in {handler.__class__}.mro()",
+                f"Expected to see HubAuthenticatedHandler in {handler.__class__}.mro(),"
+                " patching in at call time. Hub authentication is still applied.",
                 RuntimeWarning,
                 stacklevel=2,
             )
+            # patch HubAuthenticated into the instance
             handler.__class__ = type(
                 handler.__class__.__name__,
                 (HubAuthenticatedHandler, handler.__class__),
                 {},
             )
+            # patch into the class itself so this doesn't happen again for the same class
+            patch_base_handler(handler.__class__)
         return handler.get_current_user()
 
     @classmethod
@@ -683,6 +692,97 @@ def detect_base_package(App):
     return None
 
 
+def _nice_cls_repr(cls):
+    """Nice repr of classes, e.g. 'module.submod.Class'
+
+    Also accepts tuples of classes
+    """
+    return f"{cls.__module__}.{cls.__name__}"
+
+
+def patch_base_handler(BaseHandler, log=None):
+    """Patch HubAuthenticated into a base handler class
+
+    so anything inheriting from BaseHandler uses Hub authentication.
+    This works *even after* subclasses have imported and inherited from BaseHandler.
+
+    .. versionadded: 1.5
+        Made available as an importable utility
+    """
+    if log is None:
+        log = logging.getLogger()
+
+    if HubAuthenticatedHandler not in BaseHandler.__bases__:
+        new_bases = (HubAuthenticatedHandler,) + BaseHandler.__bases__
+        log.info(
+            "Patching auth into {mod}.{name}({old_bases}) -> {name}({new_bases})".format(
+                mod=BaseHandler.__module__,
+                name=BaseHandler.__name__,
+                old_bases=', '.join(
+                    _nice_cls_repr(cls) for cls in BaseHandler.__bases__
+                ),
+                new_bases=', '.join(_nice_cls_repr(cls) for cls in new_bases),
+            )
+        )
+        BaseHandler.__bases__ = new_bases
+        # We've now inserted our class as a parent of BaseHandler,
+        # but we also need to ensure BaseHandler *itself* doesn't
+        # override the public tornado API methods we have inserted.
+        # If they are defined in BaseHandler, explicitly replace them with our methods.
+        for name in ("get_current_user", "get_login_url"):
+            if name in BaseHandler.__dict__:
+                log.debug(
+                    f"Overriding {BaseHandler}.{name} with HubAuthenticatedHandler.{name}"
+                )
+                method = getattr(HubAuthenticatedHandler, name)
+                setattr(BaseHandler, name, method)
+    return BaseHandler
+
+
+def _patch_app_base_handlers(app):
+    """Patch Hub Authentication into the base handlers of an app
+
+    Patches HubAuthenticatedHandler into:
+
+    - App.base_handler_class (if defined)
+    - jupyter_server's JupyterHandler (if already imported)
+    - notebook's IPythonHandler (if already imported)
+    """
+    BaseHandler = app_base_handler = getattr(app, "base_handler_class", None)
+
+    base_handlers = []
+    if BaseHandler is not None:
+        base_handlers.append(BaseHandler)
+
+    # patch juptyer_server and notebook handlers if they have been imported
+    for base_handler_name in [
+        "jupyter_server.base.handlers.JupyterHandler",
+        "notebook.base.handlers.IPythonHandler",
+    ]:
+        modname, _ = base_handler_name.rsplit(".", 1)
+        if modname in sys.modules:
+            base_handlers.append(import_item(base_handler_name))
+
+    if not base_handlers:
+        pkg = detect_base_package(app.__class__)
+        if pkg == "jupyter_server":
+            BaseHandler = import_item("jupyter_server.base.handlers.JupyterHandler")
+        elif pkg == "notebook":
+            BaseHandler = import_item("notebook.base.handlers.IPythonHandler")
+        else:
+            raise ValueError(
+                "{}.base_handler_class must be defined".format(app.__class__.__name__)
+            )
+        base_handlers.append(BaseHandler)
+
+    # patch-in HubAuthenticatedHandler to base handler classes
+    for BaseHandler in base_handlers:
+        patch_base_handler(BaseHandler)
+
+    # return the first entry
+    return base_handlers[0]
+
+
 def make_singleuser_app(App):
     """Make and return a singleuser notebook app
 
@@ -706,37 +806,7 @@ def make_singleuser_app(App):
     # detect base classes
     LoginHandler = empty_parent_app.login_handler_class
     LogoutHandler = empty_parent_app.logout_handler_class
-    BaseHandler = getattr(empty_parent_app, "base_handler_class", None)
+    BaseHandler = _patch_app_base_handlers(empty_parent_app)
-    if BaseHandler is None:
-        pkg = detect_base_package(App)
-        if pkg == "jupyter_server":
-            BaseHandler = import_item("jupyter_server.base.handlers.JupyterHandler")
-        elif pkg == "notebook":
-            BaseHandler = import_item("notebook.base.handlers.IPythonHandler")
-        else:
-            raise ValueError(
-                "{}.base_handler_class must be defined".format(App.__name__)
-            )
-
-    # patch-in HubAuthenticatedHandler to BaseHandler,
-    # so anything inheriting from BaseHandler uses Hub authentication
-    if HubAuthenticatedHandler not in BaseHandler.__bases__:
-        new_bases = (HubAuthenticatedHandler,) + BaseHandler.__bases__
-        log.debug(
-            f"Patching {BaseHandler}{BaseHandler.__bases__} -> {BaseHandler}{new_bases}"
-        )
-        BaseHandler.__bases__ = new_bases
-        # We've now inserted our class as a parent of BaseHandler,
-        # but we also need to ensure BaseHandler *itself* doesn't
-        # override the public tornado API methods we have inserted.
-        # If they are defined in BaseHandler, explicitly replace them with our methods.
-        for name in ("get_current_user", "get_login_url"):
-            if name in BaseHandler.__dict__:
-                log.debug(
-                    f"Overriding {BaseHandler}.{name} with HubAuthenticatedHandler.{name}"
-                )
-                method = getattr(HubAuthenticatedHandler, name)
-                setattr(BaseHandler, name, method)
 
     # create Handler classes from mixins + bases
     class JupyterHubLoginHandler(JupyterHubLoginHandlerMixin, LoginHandler):
@@ -766,4 +836,11 @@ def make_singleuser_app(App):
         logout_handler_class = JupyterHubLogoutHandler
         oauth_callback_handler_class = OAuthCallbackHandler
 
+        def initialize(self, *args, **kwargs):
+            result = super().initialize(*args, **kwargs)
+            # run patch again after initialize, so extensions have already been loaded
+            # probably a no-op most of the time
+            _patch_app_base_handlers(self)
+            return result
+
     return SingleUserNotebookApp