Fix OIDC endpoints detection

Fix oidc_endpoints property to separate Databricks and Azure OIDC endpoints.

## Problem
The oidc_endpoints property incorrectly returned Azure OIDC endpoints when
ARM_CLIENT_ID was set, even for Databricks OAuth flows (like oauth-m2m)
that don't use Azure authentication. This caused Databricks M2M OAuth to
fail when users set ARM_CLIENT_ID for other purposes.

## Solution
- Created databricks_oidc_endpoints property for Databricks OIDC only
- Kept oidc_endpoints for backward compatibility (marked as deprecated)
- Updated all Databricks OAuth flows to use databricks_oidc_endpoints
- Updated Azure-specific flows to explicitly use Azure endpoints

## Tests
- Added 4 integration tests for OAuth M2M and Azure Client Secret auth
- Added 6 unit tests covering all scenarios including the bug case

Mirrors: databricks/databricks-sdk-java#657

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: hectorcast-db

databricks/sdk/config.py

@@ -486,12 +486,23 @@ def with_user_agent_extra(self, key: str, value: str) -> "Config":
         return self
 
     @property
-    def oidc_endpoints(self) -> Optional[OidcEndpoints]:
+    def databricks_oidc_endpoints(self) -> Optional[OidcEndpoints]:
+        """Get OIDC endpoints for Databricks OAuth.
+
+        This method returns the appropriate Databricks OIDC endpoints based on the host type:
+        - Unified hosts: Returns unified account-scoped endpoints
+        - Account hosts: Returns traditional account endpoints
+        - Workspace hosts: Returns workspace endpoints
+
+        Note: This method does NOT return Azure Entra ID endpoints. For Azure authentication,
+        use get_azure_entra_id_workspace_endpoints() directly.
+
+        Returns:
+            OidcEndpoints for Databricks OAuth, or None if host is not configured.
+        """
         self._fix_host_if_needed()
         if not self.host:
             return None
-        if self.is_azure and self.azure_client_id:
-            return get_azure_entra_id_workspace_endpoints(self.host)
 
         # Handle unified hosts
         if self.host_type == HostType.UNIFIED:
@@ -506,6 +517,28 @@ def oidc_endpoints(self) -> Optional[OidcEndpoints]:
         # Default to workspace endpoints
         return get_workspace_endpoints(self.host)
 
+    @property
+    def oidc_endpoints(self) -> Optional[OidcEndpoints]:
+        """[DEPRECATED] Get OIDC endpoints with automatic Azure detection (deprecated).
+
+        This method incorrectly returns Azure OIDC endpoints when azure_client_id
+        is set, even for Databricks OAuth flows that don't use Azure authentication. This caused
+        bugs where Databricks M2M OAuth would fail when ARM_CLIENT_ID was set for other purposes.
+
+        Use instead:
+        - databricks_oidc_endpoints: For Databricks OAuth (oauth-m2m, external-browser, etc.)
+        - get_azure_entra_id_workspace_endpoints(): For Azure Entra ID authentication
+
+        Returns:
+            OidcEndpoints (Azure or Databricks depending on config), or None if host is not configured.
+        """
+        self._fix_host_if_needed()
+        if not self.host:
+            return None
+        if self.is_azure and self.azure_client_id:
+            return get_azure_entra_id_workspace_endpoints(self.host)
+        return self.databricks_oidc_endpoints
+
     def debug_string(self) -> str:
         """Returns log-friendly representation of configured attributes"""
         buf = []

databricks/sdk/credentials_provider.py

@@ -20,6 +20,8 @@
 from google.auth.transport.requests import Request  # type: ignore
 from google.oauth2 import service_account  # type: ignore
 
+from databricks.sdk.oauth import get_azure_entra_id_workspace_endpoints
+
 from . import azure, oauth, oidc, oidc_token_supplier
 from .client_types import ClientType
 
@@ -218,7 +220,7 @@ def oauth_service_principal(cfg: "Config") -> Optional[CredentialsProvider]:
     """Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request,
     if /oidc/.well-known/oauth-authorization-server is available on the given host.
     """
-    oidc = cfg.oidc_endpoints
+    oidc = cfg.databricks_oidc_endpoints
     if oidc is None:
         return None
 
@@ -248,14 +250,21 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]:
         return None
 
     client_id, client_secret = None, None
+    oidc_endpoints = None
     if cfg.client_id:
         client_id = cfg.client_id
         client_secret = cfg.client_secret
+        oidc_endpoints = cfg.databricks_oidc_endpoints
     elif cfg.azure_client_id:
-        client_id = cfg.azure_client
+        client_id = cfg.azure_client_id
         client_secret = cfg.azure_client_secret
+        oidc_endpoints = get_azure_entra_id_workspace_endpoints(cfg.host)
     if not client_id:
         client_id = "databricks-cli"
+        oidc_endpoints = cfg.databricks_oidc_endpoints
+
+    if not oidc_endpoints:
+        return None
 
     scopes = cfg.get_scopes()
     if not cfg.disable_oauth_refresh_token:
@@ -264,7 +273,6 @@ def external_browser(cfg: "Config") -> Optional[CredentialsProvider]:
 
     # Load cached credentials from disk if they exist. Note that these are
     # local to the Python SDK and not reused by other SDKs.
-    oidc_endpoints = cfg.oidc_endpoints
     redirect_url = "http://localhost:8020"
     token_cache = oauth.TokenCache(
         host=cfg.host,
@@ -390,7 +398,7 @@ def oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optio
 
     token_source = oidc.DatabricksOidcTokenSource(
         host=cfg.host,
-        token_endpoint=cfg.oidc_endpoints.token_endpoint,
+        token_endpoint=cfg.databricks_oidc_endpoints.token_endpoint,
         client_id=cfg.client_id,
         account_id=cfg.account_id,
         id_token_source=id_token_source,
@@ -434,7 +442,7 @@ def _oidc_credentials_provider(
     if audience is None and cfg.client_type == ClientType.ACCOUNT:
         audience = cfg.account_id
     if audience is None and cfg.client_type != ClientType.ACCOUNT:
-        audience = cfg.oidc_endpoints.token_endpoint
+        audience = cfg.databricks_oidc_endpoints.token_endpoint
 
     # Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.
     id_token = supplier.get_oidc_token(audience)
@@ -453,7 +461,7 @@ def token_source_for(audience: str) -> oauth.TokenSource:
         return oauth.ClientCredentials(
             client_id=cfg.client_id,
             client_secret="",  # we have no (rotatable) secrets in OIDC flow
-            token_url=cfg.oidc_endpoints.token_endpoint,
+            token_url=cfg.databricks_oidc_endpoints.token_endpoint,
             endpoint_params={
                 "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
                 "subject_token": id_token,
@@ -528,7 +536,7 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
     aad_endpoint = cfg.arm_environment.active_directory_endpoint
     if not cfg.azure_tenant_id:
         # detect Azure AD Tenant ID if it's not specified directly
-        token_endpoint = cfg.oidc_endpoints.token_endpoint
+        token_endpoint = get_azure_entra_id_workspace_endpoints(cfg.host).token_endpoint
         cfg.azure_tenant_id = token_endpoint.replace(aad_endpoint, "").split("/")[0]
 
     inner = oauth.ClientCredentials(

databricks/sdk/oauth.py

@@ -692,7 +692,7 @@ def noop_credentials(_: any):
             return lambda: {}
 
         config = Config(host=host, credentials_strategy=noop_credentials)
-        oidc = config.oidc_endpoints
+        oidc = config.databricks_oidc_endpoints
         if not oidc:
             raise ValueError(f"{host} does not support OAuth")
         return OAuthClient(oidc, redirect_url, client_id, scopes, client_secret)

tests/integration/test_auth.py

@@ -269,3 +269,98 @@ def test_wif_workspace(ucacct, env_or_skip, random):
     )
 
     ws.current_user.me()
+
+
+def test_workspace_oauth_m2m_auth(w, env_or_skip):
+    env_or_skip("CLOUD_ENV")
+
+    # Get environment variables
+    host = env_or_skip("DATABRICKS_HOST")
+    client_id = env_or_skip("TEST_DATABRICKS_CLIENT_ID")
+    client_secret = env_or_skip("TEST_DATABRICKS_CLIENT_SECRET")
+
+    # Create workspace client with OAuth M2M authentication
+    ws = WorkspaceClient(
+        host=host,
+        client_id=client_id,
+        client_secret=client_secret,
+        auth_type="oauth-m2m",
+    )
+
+    # Call the "me" API
+    me = ws.current_user.me()
+
+    # Verify we got a valid response
+    assert me.user_name, "expected non-empty user_name"
+
+
+def test_workspace_azure_client_secret_auth(w, env_or_skip):
+    env_or_skip("CLOUD_ENV")
+
+    host = env_or_skip("DATABRICKS_HOST")
+    azure_client_id = env_or_skip("ARM_CLIENT_ID")
+    azure_client_secret = env_or_skip("ARM_CLIENT_SECRET")
+    azure_tenant_id = env_or_skip("ARM_TENANT_ID")
+
+    # Create workspace client with Azure client secret authentication
+    ws = WorkspaceClient(
+        host=host,
+        azure_client_id=azure_client_id,
+        azure_client_secret=azure_client_secret,
+        azure_tenant_id=azure_tenant_id,
+        auth_type="azure-client-secret",
+    )
+
+    # Call the "me" API
+    me = ws.current_user.me()
+
+    # Verify we got a valid response
+    assert me.user_name, "expected non-empty user_name"
+
+
+def test_account_oauth_m2m_auth(a, env_or_skip):
+    env_or_skip("CLOUD_ENV")
+
+    # Get environment variables
+    host = env_or_skip("DATABRICKS_HOST")
+    account_id = env_or_skip("DATABRICKS_ACCOUNT_ID")
+    client_id = env_or_skip("TEST_DATABRICKS_CLIENT_ID")
+    client_secret = env_or_skip("TEST_DATABRICKS_CLIENT_SECRET")
+
+    # Create account client with OAuth M2M authentication
+    ac = AccountClient(
+        host=host,
+        account_id=account_id,
+        client_id=client_id,
+        client_secret=client_secret,
+        auth_type="oauth-m2m",
+    )
+
+    # List service principals to verify authentication works
+    sps = ac.service_principals.list()
+    next(sps)
+
+
+def test_account_azure_client_secret_auth(a, env_or_skip):
+    env_or_skip("CLOUD_ENV")
+
+    # Get environment variables
+    host = env_or_skip("DATABRICKS_HOST")
+    account_id = env_or_skip("DATABRICKS_ACCOUNT_ID")
+    azure_client_id = env_or_skip("ARM_CLIENT_ID")
+    azure_client_secret = env_or_skip("ARM_CLIENT_SECRET")
+    azure_tenant_id = env_or_skip("ARM_TENANT_ID")
+
+    # Create account client with Azure client secret authentication
+    ac = AccountClient(
+        host=host,
+        account_id=account_id,
+        azure_client_id=azure_client_id,
+        azure_client_secret=azure_client_secret,
+        azure_tenant_id=azure_tenant_id,
+        auth_type="azure-client-secret",
+    )
+
+    # List service principals to verify authentication works
+    sps = ac.service_principals.list()
+    next(sps)

tests/test_config.py

@@ -422,6 +422,142 @@ def test_oidc_endpoints_unified_missing_ids():
     assert "Unified host requires account_id" in str(exc_info.value)
 
 
+def test_databricks_oidc_endpoints_ignores_azure_client_id(mocker, requests_mock):
+    """Test that databricks_oidc_endpoints returns Databricks endpoints even when azure_client_id is set."""
+    requests_mock.get(
+        "https://adb-123.4.azuredatabricks.net/oidc/.well-known/oauth-authorization-server",
+        json={
+            "authorization_endpoint": "https://adb-123.4.azuredatabricks.net/oidc/v1/authorize",
+            "token_endpoint": "https://adb-123.4.azuredatabricks.net/oidc/v1/token",
+        },
+    )
+
+    config = Config(
+        host="https://adb-123.4.azuredatabricks.net",
+        azure_client_id="test-azure-client-id",  # This should be ignored by databricks_oidc_endpoints
+        token="test-token",
+    )
+
+    endpoints = config.databricks_oidc_endpoints
+    assert endpoints is not None
+    assert "https://adb-123.4.azuredatabricks.net/oidc/v1/authorize" == endpoints.authorization_endpoint
+    assert "https://adb-123.4.azuredatabricks.net/oidc/v1/token" == endpoints.token_endpoint
+
+
+def test_databricks_oidc_endpoints_unified_workspace(mocker, requests_mock):
+    """Test that databricks_oidc_endpoints returns unified endpoints for workspace on unified host."""
+    requests_mock.get(
+        "https://unified.databricks.com/oidc/accounts/test-account/.well-known/oauth-authorization-server",
+        json={
+            "authorization_endpoint": "https://unified.databricks.com/oidc/accounts/test-account/v1/authorize",
+            "token_endpoint": "https://unified.databricks.com/oidc/accounts/test-account/v1/token",
+        },
+    )
+
+    config = Config(
+        host="https://unified.databricks.com",
+        workspace_id="test-workspace",
+        account_id="test-account",
+        experimental_is_unified_host=True,
+        token="test-token",
+    )
+
+    endpoints = config.databricks_oidc_endpoints
+    assert endpoints is not None
+    assert "accounts/test-account" in endpoints.authorization_endpoint
+    assert "accounts/test-account" in endpoints.token_endpoint
+
+
+def test_databricks_oidc_endpoints_account(mocker, requests_mock):
+    """Test that databricks_oidc_endpoints returns account endpoints for account hosts."""
+    requests_mock.get(
+        "https://accounts.cloud.databricks.com/oidc/accounts/test-account/.well-known/oauth-authorization-server",
+        json={
+            "authorization_endpoint": "https://accounts.cloud.databricks.com/oidc/accounts/test-account/v1/authorize",
+            "token_endpoint": "https://accounts.cloud.databricks.com/oidc/accounts/test-account/v1/token",
+        },
+    )
+
+    config = Config(
+        host="https://accounts.cloud.databricks.com",
+        account_id="test-account",
+        token="test-token",
+    )
+
+    endpoints = config.databricks_oidc_endpoints
+    assert endpoints is not None
+    assert "accounts/test-account" in endpoints.authorization_endpoint
+    assert "accounts/test-account" in endpoints.token_endpoint
+
+
+def test_databricks_oidc_endpoints_workspace(mocker, requests_mock):
+    """Test that databricks_oidc_endpoints returns workspace endpoints for workspace hosts."""
+    requests_mock.get(
+        "https://test-workspace.cloud.databricks.com/oidc/.well-known/oauth-authorization-server",
+        json={
+            "authorization_endpoint": "https://test-workspace.cloud.databricks.com/oidc/v1/authorize",
+            "token_endpoint": "https://test-workspace.cloud.databricks.com/oidc/v1/token",
+        },
+    )
+
+    config = Config(
+        host="https://test-workspace.cloud.databricks.com",
+        token="test-token",
+    )
+
+    endpoints = config.databricks_oidc_endpoints
+    assert endpoints is not None
+    assert "https://test-workspace.cloud.databricks.com/oidc/v1/authorize" == endpoints.authorization_endpoint
+    assert "https://test-workspace.cloud.databricks.com/oidc/v1/token" == endpoints.token_endpoint
+
+
+def test_oidc_endpoints_returns_azure_when_azure_client_id_set(mocker):
+    """Test that deprecated oidc_endpoints returns Azure endpoints when azure_client_id is set on Azure.
+
+    This tests the deprecated behavior that is maintained for backward compatibility.
+    """
+    # Mock the Azure endpoint detection
+    mocker.patch(
+        "databricks.sdk.config.get_azure_entra_id_workspace_endpoints",
+        return_value=mocker.Mock(
+            authorization_endpoint="https://login.microsoftonline.com/tenant-id/oauth2/v2.0/authorize",
+            token_endpoint="https://login.microsoftonline.com/tenant-id/oauth2/v2.0/token",
+        ),
+    )
+
+    config = Config(
+        host="https://adb-123.4.azuredatabricks.net",
+        azure_client_id="test-azure-client-id",
+        token="test-token",
+    )
+
+    endpoints = config.oidc_endpoints
+    assert endpoints is not None
+    assert "login.microsoftonline.com" in endpoints.authorization_endpoint
+    assert "login.microsoftonline.com" in endpoints.token_endpoint
+
+
+def test_oidc_endpoints_falls_back_to_databricks_when_no_azure_client_id(mocker, requests_mock):
+    """Test that deprecated oidc_endpoints falls back to Databricks endpoints when azure_client_id is not set."""
+    requests_mock.get(
+        "https://adb-123.4.azuredatabricks.net/oidc/.well-known/oauth-authorization-server",
+        json={
+            "authorization_endpoint": "https://adb-123.4.azuredatabricks.net/oidc/v1/authorize",
+            "token_endpoint": "https://adb-123.4.azuredatabricks.net/oidc/v1/token",
+        },
+    )
+
+    config = Config(
+        host="https://adb-123.4.azuredatabricks.net",
+        token="test-token",
+    )
+
+    endpoints = config.oidc_endpoints
+    assert endpoints is not None
+    assert "https://adb-123.4.azuredatabricks.net/oidc/v1/authorize" == endpoints.authorization_endpoint
+    assert "https://adb-123.4.azuredatabricks.net/oidc/v1/token" == endpoints.token_endpoint
+
+
 def test_workspace_org_id_header_on_unified_host(requests_mock):
     """Test that X-Databricks-Org-Id header is added for workspace clients on unified hosts."""