Skip to content

Commit

Permalink
Reduce redundant warning messages in azure-cli credential provider (#410
Browse files Browse the repository at this point in the history
)

## Changes
This warning is a bit noisy: in the happy path, we print it multiple
times, and even if we don't use Azure CLI auth, we may print it anyways.
Also, it will be printed always when using account-level authentication.

This PR reduces the noise for this warning by only printing it when we
are certainly going to use Azure CLI auth, authenticating to a
workspace, and can't get the subscription ID (e.g. because it is not
configured).

## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] Ran locally: with accounts client, no warning is printed, and with
workspace client, only one warning is printed.
  • Loading branch information
mgyucht authored Oct 23, 2023
1 parent d4c160f commit f29ad4a
Showing 1 changed file with 9 additions and 8 deletions.
17 changes: 9 additions & 8 deletions databricks/sdk/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -316,7 +316,7 @@ def __init__(self, resource: str, subscription: str = ""):

@staticmethod
def for_resource(cfg: 'Config', resource: str) -> 'AzureCliTokenSource':
subscription = AzureCliTokenSource._get_subscription(cfg)
subscription = AzureCliTokenSource.get_subscription(cfg)
if subscription != "":
token = AzureCliTokenSource(resource, subscription)
try:
Expand All @@ -327,19 +327,15 @@ def for_resource(cfg: 'Config', resource: str) -> 'AzureCliTokenSource':
return token
except OSError:
logger.warning("Failed to get token for subscription. Using resource only token.")
else:
logger.warning(
"azure_workspace_resource_id field not provided. " +
"It is recommended to specify this field in the Databricks configuration to avoid authentication errors."
)

token = AzureCliTokenSource(resource)
token.token()
return token

@staticmethod
def _get_subscription(cfg: 'Config') -> str:
def get_subscription(cfg: 'Config') -> str:
resource = cfg.azure_workspace_resource_id
if resource == None or resource == "":
if resource is None or resource == "":
return ""
components = resource.split('/')
if len(components) < 3:
Expand Down Expand Up @@ -368,6 +364,11 @@ def azure_cli(cfg: 'Config') -> Optional[HeaderFactory]:

_ensure_host_present(cfg, lambda resource: AzureCliTokenSource.for_resource(cfg, resource))
logger.info("Using Azure CLI authentication with AAD tokens")
if not cfg.is_account_client and AzureCliTokenSource.get_subscription(cfg) == "":
logger.warning(
"azure_workspace_resource_id field not provided. "
"It is recommended to specify this field in the Databricks configuration to avoid authentication errors."
)

def inner() -> Dict[str, str]:
token = token_source.token()
Expand Down

0 comments on commit f29ad4a

Please sign in to comment.