From cef68f8a9af41d36c22557fedba42263d94a5ed4 Mon Sep 17 00:00:00 2001
From: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
Date: Wed, 11 Oct 2023 15:38:31 +0200
Subject: [PATCH] fix: Breaking change in MachineAuthProvider constructor
 (#25532)

---
 superset/config.py             |  1 -
 superset/utils/machine_auth.py | 17 +++++++++--------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index 7463d3083f580..63b2e0dba13b7 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1355,7 +1355,6 @@ def EMAIL_HEADER_MUTATOR(  # pylint: disable=invalid-name,unused-argument
 # 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"}
diff --git a/superset/utils/machine_auth.py b/superset/utils/machine_auth.py
index 6cd1c0ba7449e..2382f9d727cfa 100644
--- a/superset/utils/machine_auth.py
+++ b/superset/utils/machine_auth.py
@@ -43,15 +43,16 @@
 class MachineAuthProvider:
     def __init__(
         self,
-        auth_webdriver_func_override: Callable[[WebDriver, User], WebDriver],
-        auth_context_func_override: Callable[[BrowserContext, User], BrowserContext],
+        auth_webdriver_func_override: Callable[
+            [WebDriver | BrowserContext, User], WebDriver | BrowserContext
+        ]
+        | None = None,
     ):
         # This is here in order to allow for the authenticate_webdriver
         # or authenticate_browser_context (if PLAYWRIGHT_REPORTS_AND_THUMBNAILS is
         # enabled) func to be overridden via config, as opposed to the entire
         # provider implementation
         self._auth_webdriver_func_override = auth_webdriver_func_override
-        self._auth_context_func_override = auth_context_func_override
 
     def authenticate_webdriver(
         self,
@@ -63,7 +64,7 @@ def authenticate_webdriver(
         :return: The WebDriver passed in (fluent)
         """
         # Short-circuit this method if we have an override configured
-        if self._auth_webdriver_func_override:  # type: ignore
+        if self._auth_webdriver_func_override:
             return self._auth_webdriver_func_override(driver, user)
 
         # Setting cookies requires doing a request first
@@ -82,8 +83,8 @@ def authenticate_browser_context(
         user: User,
     ) -> BrowserContext:
         # Short-circuit this method if we have an override configured
-        if self._auth_context_func_override:  # type: ignore
-            return self._auth_context_func_override(browser_context, user)
+        if self._auth_webdriver_func_override:
+            return self._auth_webdriver_func_override(browser_context, user)
 
         url = urlparse(current_app.config["WEBDRIVER_BASEURL"])
 
@@ -145,12 +146,12 @@ def get_auth_cookies(user: User) -> dict[str, str]:
 
 class MachineAuthProviderFactory:
     def __init__(self) -> None:
-        self._auth_provider = None
+        self._auth_provider: MachineAuthProvider | None = 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["BROWSER_CONTEXT_AUTH_FUNC"])
+        )(app.config["WEBDRIVER_AUTH_FUNC"])
 
     @property
     def instance(self) -> MachineAuthProvider: