address comment

Author: madhav-db

src/databricks/sql/auth/auth.py

@@ -8,13 +8,13 @@
     AzureServicePrincipalCredentialProvider,
 )
 from databricks.sql.auth.common import AuthType, ClientContext
-from databricks.sql.auth.token_federation import TokenFederationProvider, ExternalTokenProvider
+from databricks.sql.auth.token_federation import TokenFederationProvider
 
 
 def get_auth_provider(cfg: ClientContext, http_client):
     # Determine the base auth provider
     base_provider = None
-    
+
     if cfg.credentials_provider:
         base_provider = ExternalAuthProvider(cfg.credentials_provider)
     elif cfg.auth_type == AuthType.AZURE_SP_M2M.value:
@@ -64,16 +64,16 @@ def get_auth_provider(cfg: ClientContext, http_client):
             )
         else:
             raise RuntimeError("No valid authentication settings!")
-    
-    # Wrap with token federation if enabled
-    if cfg.enable_token_federation and base_provider:
+
+    # Always wrap with token federation (falls back gracefully if not needed)
+    if base_provider:
         return TokenFederationProvider(
             hostname=cfg.hostname,
             external_provider=base_provider,
             http_client=http_client,
             identity_federation_client_id=cfg.identity_federation_client_id,
         )
-    
+
     return base_provider
 
 
@@ -129,8 +129,6 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs)
         else redirect_port_range,
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
-        # Token federation parameters
-        enable_token_federation=kwargs.get("enable_token_federation", False),
         identity_federation_client_id=kwargs.get("identity_federation_client_id"),
     )
     return get_auth_provider(cfg, http_client)

src/databricks/sql/auth/auth_utils.py

@@ -0,0 +1,117 @@
+import logging
+import jwt
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Tuple
+from urllib.parse import urlparse
+
+logger = logging.getLogger(__name__)
+
+
+def parse_hostname(hostname: str) -> str:
+    """
+    Normalize the hostname to include scheme and trailing slash.
+
+    Args:
+        hostname: The hostname to normalize
+
+    Returns:
+        Normalized hostname with scheme and trailing slash
+    """
+    if not hostname.startswith("http://") and not hostname.startswith("https://"):
+        hostname = f"https://{hostname}"
+    if not hostname.endswith("/"):
+        hostname = f"{hostname}/"
+    return hostname
+
+
+def decode_token(access_token: str) -> Optional[Dict]:
+    """
+    Decode a JWT token without verification to extract claims.
+
+    Args:
+        access_token: The JWT access token to decode
+
+    Returns:
+        Decoded token claims or None if decoding fails
+    """
+    try:
+        return jwt.decode(access_token, options={"verify_signature": False})
+    except Exception as e:
+        logger.debug("Failed to decode JWT token: %s", e)
+        return None
+
+
+def is_same_host(url1: str, url2: str) -> bool:
+    """
+    Check if two URLs have the same host.
+
+    Args:
+        url1: First URL
+        url2: Second URL
+
+    Returns:
+        True if hosts are the same, False otherwise
+    """
+    try:
+        host1 = urlparse(url1).netloc
+        host2 = urlparse(url2).netloc
+        # Handle port differences (e.g., example.com vs example.com:443)
+        host1_without_port = host1.split(":")[0]
+        host2_without_port = host2.split(":")[0]
+        return host1_without_port == host2_without_port
+    except Exception as e:
+        logger.debug("Failed to parse URLs: %s", e)
+        return False
+
+
+class Token:
+    """
+    Represents an OAuth token with expiration management.
+    """
+
+    def __init__(self, access_token: str, token_type: str = "Bearer"):
+        """
+        Initialize a token.
+
+        Args:
+            access_token: The access token string
+            token_type: The token type (default: Bearer)
+        """
+        self.access_token = access_token
+        self.token_type = token_type
+        self.expiry_time = self._calculate_expiry()
+
+    def _calculate_expiry(self) -> datetime:
+        """
+        Calculate the token expiry time from JWT claims.
+
+        Returns:
+            The token expiry datetime
+        """
+        decoded = decode_token(self.access_token)
+        if decoded and "exp" in decoded:
+            # Use JWT exp claim with 1 minute buffer
+            return datetime.fromtimestamp(decoded["exp"]) - timedelta(minutes=1)
+        # Default to 1 hour if no expiry info
+        return datetime.now() + timedelta(hours=1)
+
+    def is_expired(self) -> bool:
+        """
+        Check if the token is expired.
+
+        Returns:
+            True if token is expired, False otherwise
+        """
+        return datetime.now() >= self.expiry_time
+
+    def to_dict(self) -> Dict[str, str]:
+        """
+        Convert token to dictionary format.
+
+        Returns:
+            Dictionary with access_token and token_type
+        """
+        return {
+            "access_token": self.access_token,
+            "token_type": self.token_type,
+        }

src/databricks/sql/auth/common.py

@@ -37,8 +37,6 @@ def __init__(
         tls_client_cert_file: Optional[str] = None,
         oauth_persistence=None,
         credentials_provider=None,
-        # Token federation parameters
-        enable_token_federation: bool = False,
         identity_federation_client_id: Optional[str] = None,
         # HTTP client configuration parameters
         ssl_options=None,  # SSLOptions type
@@ -68,8 +66,6 @@ def __init__(
         self.tls_client_cert_file = tls_client_cert_file
         self.oauth_persistence = oauth_persistence
         self.credentials_provider = credentials_provider
-        # Token federation
-        self.enable_token_federation = enable_token_federation
         self.identity_federation_client_id = identity_federation_client_id
 
         # HTTP client configuration