|
13 | 13 | from . import useragent |
14 | 14 | from ._base_client import _fix_host_if_needed |
15 | 15 | from .clock import Clock, RealClock |
16 | | -from .credentials_provider import CredentialsStrategy, DefaultCredentials |
| 16 | +from .credentials_provider import CredentialsStrategy, DefaultCredentials, OAuthCredentialsProvider |
17 | 17 | from .environments import (ALL_ENVS, AzureEnvironment, Cloud, |
18 | 18 | DatabricksEnvironment, get_environment_for_hostname) |
19 | 19 | from .oauth import (OidcEndpoints, Token, get_account_endpoints, |
@@ -200,7 +200,19 @@ def __init__( |
200 | 200 | raise ValueError(message) from e |
201 | 201 |
|
202 | 202 | def oauth_token(self) -> Token: |
203 | | - return self._credentials_strategy.oauth_token(self) |
| 203 | + """Returns the OAuth token from the current credential provider. |
| 204 | + |
| 205 | + This method only works when using OAuth-based authentication methods. |
| 206 | + If the current credential provider is an OAuthCredentialsProvider, it reuses |
| 207 | + the existing provider. Otherwise, it raises a ValueError indicating that |
| 208 | + OAuth tokens are not available for the current authentication method. |
| 209 | + """ |
| 210 | + # Check if the current header factory is an OAuthCredentialsProvider |
| 211 | + if isinstance(self._header_factory, OAuthCredentialsProvider): |
| 212 | + return self._header_factory.oauth_token() |
| 213 | + # Raise an error for non-OAuth authentication methods |
| 214 | + raise ValueError(f"OAuth tokens are not available for {self.auth_type} authentication. " |
| 215 | + f"Use an OAuth-based authentication method to access OAuth tokens.") |
204 | 216 |
|
205 | 217 | def wrap_debug_info(self, message: str) -> str: |
206 | 218 | debug_string = self.debug_string() |
|
0 commit comments