From 8be1f12e7eefa45f41a4c48b98a4dd8b08408917 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 11:49:31 +0000 Subject: [PATCH 1/6] Fix cache re-creation --- databricks/sdk/config.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 487527be7..6ae18cf8e 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -13,7 +13,7 @@ from . import useragent from ._base_client import _fix_host_if_needed from .clock import Clock, RealClock -from .credentials_provider import CredentialsStrategy, DefaultCredentials +from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider from .environments import (ALL_ENVS, AzureEnvironment, Cloud, DatabricksEnvironment, get_environment_for_hostname) from .oauth import (OidcEndpoints, Token, get_account_endpoints, @@ -200,7 +200,19 @@ def __init__( raise ValueError(message) from e def oauth_token(self) -> Token: - return self._credentials_strategy.oauth_token(self) + """Returns the OAuth token from the current credential provider. + + This method only works when using OAuth-based authentication methods. + If the current credential provider is an OAuthCredentialsProvider, it reuses + the existing provider. Otherwise, it raises a ValueError indicating that + OAuth tokens are not available for the current authentication method. + """ + # Check if the current header factory is an OAuthCredentialsProvider + if isinstance(self._header_factory, OAuthCredentialsProvider): + return self._header_factory.oauth_token() + # Raise an error for non-OAuth authentication methods + raise ValueError(f"OAuth tokens are not available for {self.auth_type} authentication. " + f"Use an OAuth-based authentication method to access OAuth tokens.") def wrap_debug_info(self, message: str) -> str: debug_string = self.debug_string() From 4faa7cabc3dfdab967110a663b65af5c43e0d5a2 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 11:56:30 +0000 Subject: [PATCH 2/6] Test + fmt --- databricks/sdk/config.py | 11 ++++--- tests/test_config.py | 71 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 6ae18cf8e..cc9272680 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -13,7 +13,8 @@ from . import useragent from ._base_client import _fix_host_if_needed from .clock import Clock, RealClock -from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider +from .credentials_provider import (CredentialsStrategy, DefaultCredentials, + OAuthCredentialsProvider) from .environments import (ALL_ENVS, AzureEnvironment, Cloud, DatabricksEnvironment, get_environment_for_hostname) from .oauth import (OidcEndpoints, Token, get_account_endpoints, @@ -201,7 +202,7 @@ def __init__( def oauth_token(self) -> Token: """Returns the OAuth token from the current credential provider. - + This method only works when using OAuth-based authentication methods. If the current credential provider is an OAuthCredentialsProvider, it reuses the existing provider. Otherwise, it raises a ValueError indicating that @@ -211,8 +212,10 @@ def oauth_token(self) -> Token: if isinstance(self._header_factory, OAuthCredentialsProvider): return self._header_factory.oauth_token() # Raise an error for non-OAuth authentication methods - raise ValueError(f"OAuth tokens are not available for {self.auth_type} authentication. " - f"Use an OAuth-based authentication method to access OAuth tokens.") + raise ValueError( + f"OAuth tokens are not available for {self.auth_type} authentication. " + f"Use an OAuth-based authentication method to access OAuth tokens." + ) def wrap_debug_info(self, message: str) -> str: debug_string = self.debug_string() diff --git a/tests/test_config.py b/tests/test_config.py index b023123af..a5a986a79 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -189,3 +189,74 @@ def test_load_azure_tenant_id_happy_path(requests_mock, monkeypatch): cfg = Config(host="https://abc123.azuredatabricks.net") assert cfg.azure_tenant_id == "tenant-id" assert mock.called_once + + +def test_oauth_token_with_pat_auth(): + """Test that oauth_token() raises an error for PAT authentication.""" + config = Config(host="https://test.databricks.com", token="dapi1234567890abcdef") + + with pytest.raises(ValueError) as exc_info: + config.oauth_token() + + assert "OAuth tokens are not available for pat authentication" in str(exc_info.value) + + +def test_oauth_token_with_basic_auth(): + """Test that oauth_token() raises an error for basic authentication.""" + config = Config(host="https://test.databricks.com", username="testuser", password="testpass") + + with pytest.raises(ValueError) as exc_info: + config.oauth_token() + + assert "OAuth tokens are not available for basic authentication" in str(exc_info.value) + + +def test_oauth_token_with_oauth_provider(mocker): + """Test that oauth_token() works correctly for OAuth authentication.""" + from databricks.sdk.credentials_provider import OAuthCredentialsProvider + from databricks.sdk.oauth import Token + + # Create a mock OAuth token + mock_token = Token(access_token="mock_access_token", token_type="Bearer", refresh_token="mock_refresh_token") + + # Create a mock OAuth provider + mock_oauth_provider = mocker.Mock(spec=OAuthCredentialsProvider) + mock_oauth_provider.oauth_token.return_value = mock_token + + # Create config with mocked header factory + config = Config(host="https://test.databricks.com", client_id="test-client-id", client_secret="test-client-secret") + + # Replace the header factory with our mock + config._header_factory = mock_oauth_provider + + # Test that oauth_token() works and returns the expected token + token = config.oauth_token() + assert token == mock_token + mock_oauth_provider.oauth_token.assert_called_once() + + +def test_oauth_token_reuses_existing_provider(mocker): + """Test that oauth_token() reuses the existing OAuthCredentialsProvider.""" + from databricks.sdk.credentials_provider import OAuthCredentialsProvider + from databricks.sdk.oauth import Token + + # Create a mock OAuth token + mock_token = Token(access_token="mock_access_token", token_type="Bearer", refresh_token="mock_refresh_token") + + # Create a mock OAuth provider + mock_oauth_provider = mocker.Mock(spec=OAuthCredentialsProvider) + mock_oauth_provider.oauth_token.return_value = mock_token + + # Create config with mocked header factory + config = Config(host="https://test.databricks.com", client_id="test-client-id", client_secret="test-client-secret") + + # Replace the header factory with our mock + config._header_factory = mock_oauth_provider + + # Call oauth_token() multiple times to verify reuse + token1 = config.oauth_token() + token2 = config.oauth_token() + + # Both calls should work and use the same provider instance + assert token1 == token2 == mock_token + assert mock_oauth_provider.oauth_token.call_count == 2 From 5907eb68cd87d5627c7550832146a1bb84f2e2f0 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 11:59:11 +0000 Subject: [PATCH 3/6] Add changelog --- NEXT_CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index a61cfb3a8..185bb9a33 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,6 +6,8 @@ ### Bug Fixes +* Fix `Config.oauth_token()` to avoid re-creating a new `CredentialsProvider` at each call. This fix indirectly makes `oauth_token()` benefit from the internal caching mechanism of some providers. + ### Documentation ### Internal Changes From 9d37096ed08b674339c06d9a5dd9d6412c0d7c86 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 12:00:20 +0000 Subject: [PATCH 4/6] Remove pointless comments --- databricks/sdk/config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index cc9272680..4cfd8b4f9 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -208,10 +208,8 @@ def oauth_token(self) -> Token: the existing provider. Otherwise, it raises a ValueError indicating that OAuth tokens are not available for the current authentication method. """ - # Check if the current header factory is an OAuthCredentialsProvider if isinstance(self._header_factory, OAuthCredentialsProvider): return self._header_factory.oauth_token() - # Raise an error for non-OAuth authentication methods raise ValueError( f"OAuth tokens are not available for {self.auth_type} authentication. " f"Use an OAuth-based authentication method to access OAuth tokens." From 4c07ad22132fdc7634a0b64c083d03627ddee217 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 13:20:00 +0000 Subject: [PATCH 5/6] Fix tests --- tests/test_config.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index a5a986a79..e2eb9f6dd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -223,8 +223,11 @@ def test_oauth_token_with_oauth_provider(mocker): mock_oauth_provider = mocker.Mock(spec=OAuthCredentialsProvider) mock_oauth_provider.oauth_token.return_value = mock_token - # Create config with mocked header factory - config = Config(host="https://test.databricks.com", client_id="test-client-id", client_secret="test-client-secret") + # Create config with noop credentials to avoid network calls + config = Config( + host="https://test.databricks.com", + credentials_strategy=noop_credentials + ) # Replace the header factory with our mock config._header_factory = mock_oauth_provider @@ -247,8 +250,11 @@ def test_oauth_token_reuses_existing_provider(mocker): mock_oauth_provider = mocker.Mock(spec=OAuthCredentialsProvider) mock_oauth_provider.oauth_token.return_value = mock_token - # Create config with mocked header factory - config = Config(host="https://test.databricks.com", client_id="test-client-id", client_secret="test-client-secret") + # Create config with noop credentials to avoid network calls + config = Config( + host="https://test.databricks.com", + credentials_strategy=noop_credentials + ) # Replace the header factory with our mock config._header_factory = mock_oauth_provider From 23143b6989e8ec07712bb7add214724fb59cbaa7 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 31 Jul 2025 13:27:28 +0000 Subject: [PATCH 6/6] fmt --- tests/test_config.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index e2eb9f6dd..59fbf8712 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -224,10 +224,7 @@ def test_oauth_token_with_oauth_provider(mocker): mock_oauth_provider.oauth_token.return_value = mock_token # Create config with noop credentials to avoid network calls - config = Config( - host="https://test.databricks.com", - credentials_strategy=noop_credentials - ) + config = Config(host="https://test.databricks.com", credentials_strategy=noop_credentials) # Replace the header factory with our mock config._header_factory = mock_oauth_provider @@ -251,10 +248,7 @@ def test_oauth_token_reuses_existing_provider(mocker): mock_oauth_provider.oauth_token.return_value = mock_token # Create config with noop credentials to avoid network calls - config = Config( - host="https://test.databricks.com", - credentials_strategy=noop_credentials - ) + config = Config(host="https://test.databricks.com", credentials_strategy=noop_credentials) # Replace the header factory with our mock config._header_factory = mock_oauth_provider