feat (oauth): add custom oauth manager

vaimdev · vaimdev · commit 5b8f72064932 · 2024-01-11T23:13:39.000+08:00
diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py
@@ -33,6 +33,7 @@ def __init__(
         tls_client_cert_file: str = None,
         oauth_persistence=None,
         credentials_provider=None,
+        custom_oauth_manager=None,
     ):
         self.hostname = hostname
         self.username = username
@@ -46,11 +47,13 @@ def __init__(
         self.tls_client_cert_file = tls_client_cert_file
         self.oauth_persistence = oauth_persistence
         self.credentials_provider = credentials_provider
+        self.custom_oauth_manager=custom_oauth_manager
 
 
 def get_auth_provider(cfg: ClientContext):
     if cfg.credentials_provider:
         return ExternalAuthProvider(cfg.credentials_provider)
+
     if cfg.auth_type == AuthType.DATABRICKS_OAUTH.value:
         assert cfg.oauth_redirect_port_range is not None
         assert cfg.oauth_client_id is not None
@@ -62,6 +65,7 @@ def get_auth_provider(cfg: ClientContext):
             cfg.oauth_redirect_port_range,
             cfg.oauth_client_id,
             cfg.oauth_scopes,
+            cfg.custom_oauth_manager,
         )
     elif cfg.access_token is not None:
         return AccessTokenAuthProvider(cfg.access_token)
@@ -112,5 +116,8 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
         else redirect_port_range,
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
+        # customization start
+        custom_oauth_manager=kwargs.get("custom_oauth_manager"),
+        # customization end
     )
     return get_auth_provider(cfg)
diff --git a/src/databricks/sql/auth/authenticators.py b/src/databricks/sql/auth/authenticators.py
@@ -69,6 +69,7 @@ def __init__(
         redirect_port_range: List[int],
         client_id: str,
         scopes: List[str],
+        custom_oauth_manager,
     ):
         try:
             cloud_type = infer_cloud_from_host(hostname)
@@ -84,7 +85,10 @@ def __init__(
             # Convert to the corresponding scopes in the corresponding IdP
             cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
 
-            self.oauth_manager = OAuthManager(
+            if not custom_oauth_manager:
+                custom_oauth_manager = OAuthManager
+
+            self.oauth_manager = custom_oauth_manager(
                 port_range=redirect_port_range,
                 client_id=client_id,
                 idp_endpoint=idp_endpoint,
