Separate webdriver and browsercontext auth overrides

superset/config.py

@@ -1351,9 +1351,11 @@ def EMAIL_HEADER_MUTATOR(  # pylint: disable=invalid-name,unused-argument
     "pixel_density": 1,
 }
 
-# An optional override to the default auth hook used to provide auth to the
-# offline webdriver
+# An optional override to the default auth hook used to provide auth to the offline
+# webdriver (when using Selenium) or browser context (when using Playwright - see
+# PLAYWRIGHT_REPORTS_AND_THUMBNAILS feature flag)
 WEBDRIVER_AUTH_FUNC = None
+BROWSER_CONTEXT_AUTH_FUNC = None
 
 # Any config options to be passed as-is to the webdriver
 WEBDRIVER_CONFIGURATION: dict[Any, Any] = {"service_log_path": "/dev/null"}

superset/utils/machine_auth.py

@@ -42,13 +42,15 @@
 class MachineAuthProvider:
     def __init__(
         self,
-        auth_webdriver_func_override: Callable[
-            [WebDriver | BrowserContext, User], WebDriver | BrowserContext
+        auth_webdriver_func_override: Callable[[WebDriver, User], WebDriver],
+        auth_browser_context_func_override: Callable[
+            [BrowserContext, User], BrowserContext
         ],
     ):
         # This is here in order to allow for the authenticate_webdriver func to be
         # overridden via config, as opposed to the entire provider implementation
         self._auth_webdriver_func_override = auth_webdriver_func_override
+        self._auth_browser_context_func_override = auth_browser_context_func_override
 
     def authenticate_webdriver(
         self,
@@ -79,8 +81,8 @@ def authenticate_browser_context(
         user: User,
     ) -> BrowserContext:
         # Short-circuit this method if we have an override configured
-        if self._auth_webdriver_func_override:  # type: ignore
-            return self._auth_webdriver_func_override(browser_context, user)
+        if self._auth_browser_context_func_override:  # type: ignore
+            return self._auth_browser_context_func_override(browser_context, user)
 
         # Setting cookies requires doing a request first
         page = browser_context.new_page()
@@ -145,7 +147,7 @@ def __init__(self) -> None:
     def init_app(self, app: Flask) -> None:
         self._auth_provider = load_class_from_name(
             app.config["MACHINE_AUTH_PROVIDER_CLASS"]
-        )(app.config["WEBDRIVER_AUTH_FUNC"])
+        )(app.config["WEBDRIVER_AUTH_FUNC"], app.config["BROWSER_CONTEXT_AUTH_FUNC"])
 
     @property
     def instance(self) -> MachineAuthProvider: