diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index d98c20035..f22f5d5b5 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -13,5 +13,6 @@ ### Breaking Changes ### Internal Changes +* Pass `--force-refresh` to Databricks CLI `auth token` command so the SDK always receives a freshly minted token instead of a potentially stale one from the CLI's internal cache. ### API Changes diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 090c54edd..1d46a72c1 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -1012,6 +1012,9 @@ def _validate_token_scopes(self, token: oauth.Token): # --profile support added in CLI v0.207.1: https://github.com/databricks/cli/pull/855 _CLI_VERSION_FOR_PROFILE = CliVersion(0, 207, 1) + # --force-refresh support added in CLI v0.296.0: https://github.com/databricks/cli/pull/4767 + _CLI_VERSION_FOR_FORCE_REFRESH = CliVersion(0, 296, 0) + @staticmethod def _parse_cli_version(output: str) -> CliVersion: """Parse the JSON output of `databricks version --output json`. @@ -1101,7 +1104,34 @@ def _resolve_cli_command(cli_path: str, cfg: "Config") -> List[str]: @staticmethod def _build_cli_command(cli_path: str, cfg: "Config", version: CliVersion) -> List[str]: - """Build the `auth token` command. + """Build the full CLI command, including capability-gated flags. + + Delegates the profile/host decision to _build_core_cli_command and + appends --force-refresh when supported. + """ + cmd = DatabricksCliTokenSource._build_core_cli_command(cli_path, cfg, version) + if version >= DatabricksCliTokenSource._CLI_VERSION_FOR_FORCE_REFRESH: + cmd.append("--force-refresh") + elif version == CliVersion() or version.is_default_dev_build: + # Detection failed or no version metadata — we can't prove the CLI + # lacks --force-refresh, just failed to confirm it. Upstream has + # already logged the underlying cause, so use softer wording. + logger.warning( + f"Could not confirm --force-refresh support for Databricks CLI {version} " + f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_FORCE_REFRESH}). " + "The CLI's token cache may provide stale tokens." + ) + else: + logger.warning( + f"Databricks CLI {version} does not support --force-refresh " + f"(requires >= {DatabricksCliTokenSource._CLI_VERSION_FOR_FORCE_REFRESH}). " + "The CLI's token cache may provide stale tokens." + ) + return cmd + + @staticmethod + def _build_core_cli_command(cli_path: str, cfg: "Config", version: CliVersion) -> List[str]: + """Build the base `auth token` command without capability-gated flags. Falls back to --host when --profile is either not configured or not supported by the installed CLI. Raises IOError describing which diff --git a/tests/test_credentials_provider.py b/tests/test_credentials_provider.py index eb5c0c9dc..2536ee61b 100644 --- a/tests/test_credentials_provider.py +++ b/tests/test_credentials_provider.py @@ -451,6 +451,48 @@ def _make_cfg(*, profile=None, host=None, account_id=None): _CV(0, 0, 0), [_CLI, "auth", "token", "--host", _HOST], ), + ( + "host with force-refresh", + _make_cfg(host=_HOST), + _CV(0, 296, 0), + [_CLI, "auth", "token", "--host", _HOST, "--force-refresh"], + ), + ( + "account host with force-refresh", + _make_cfg(host=_ACCT_HOST, account_id="acct-123"), + _CV(0, 296, 0), + [_CLI, "auth", "token", "--host", _ACCT_HOST, "--account-id", "acct-123", "--force-refresh"], + ), + ( + "profile with force-refresh", + _make_cfg(profile="my-profile", host=_HOST), + _CV(0, 296, 0), + [_CLI, "auth", "token", "--profile", "my-profile", "--force-refresh"], + ), + ( + "profile supports profile but not force-refresh", + _make_cfg(profile="my-profile", host=_HOST), + _CV(0, 207, 1), + [_CLI, "auth", "token", "--profile", "my-profile"], + ), + ( + "profile-only with force-refresh", + _make_cfg(profile="my-profile"), + _CV(0, 296, 0), + [_CLI, "auth", "token", "--profile", "my-profile", "--force-refresh"], + ), + ( + "unknown version, host only, no force-refresh", + _make_cfg(host=_HOST), + _CV(), + [_CLI, "auth", "token", "--host", _HOST], + ), + ( + "dev-build version, host only, no force-refresh", + _make_cfg(host=_HOST), + _CV(0, 0, 0), + [_CLI, "auth", "token", "--host", _HOST], + ), ], ) def test_build_cli_command(name, cfg, version, expected): @@ -592,6 +634,68 @@ def test_resolve_cli_command_new_cli_uses_profile(mocker): assert cmd == [_CLI, "auth", "token", "--profile", "my-profile"] +def test_build_cli_command_force_refresh_unsupported_logs_warning(caplog): + import logging + + cfg = _make_cfg(host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk.credentials_provider"): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, _CV(0, 295, 0)) + assert any( + "does not support --force-refresh" in rec.message and rec.levelname == "WARNING" for rec in caplog.records + ) + + +@pytest.mark.parametrize( + "version", + [ + # Detection failed: we don't actually know the CLI lacks --force-refresh. + _CV(), + # Default dev build: no version metadata injected, same story. + _CV(0, 0, 0), + ], +) +def test_build_cli_command_unconfirmed_force_refresh_softens_warning(caplog, version): + import logging + + cfg = _make_cfg(host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk.credentials_provider"): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, version) + # Softer phrasing for states where --force-refresh support wasn't proven absent. + assert any( + "Could not confirm --force-refresh support" in rec.message and rec.levelname == "WARNING" + for rec in caplog.records + ) + assert not any("does not support --force-refresh" in rec.message for rec in caplog.records) + + +def test_build_cli_command_force_refresh_supported_no_warning(caplog): + import logging + + cfg = _make_cfg(host=_HOST) + with caplog.at_level(logging.WARNING, logger="databricks.sdk.credentials_provider"): + credentials_provider.DatabricksCliTokenSource._build_cli_command(_CLI, cfg, _CV(0, 296, 0)) + # No --force-refresh-related warning when the flag is supported. + assert not any("--force-refresh" in rec.message for rec in caplog.records) + + +def test_resolve_cli_command_malformed_version_json_falls_back(mocker, caplog): + # Pin the integrated path: `databricks version --output json` succeeds but + # emits unparseable JSON. _parse_cli_version returns CliVersion() with only + # DEBUG logging, _resolve_cli_command falls back to --host, and no WARNING + # about --profile fires because the unknown sentinel takes the softened branch. + import logging + + _stub_version_output(mocker, "{not valid json") + cfg = _make_cfg(profile="my-profile", host=_HOST) + with caplog.at_level(logging.DEBUG, logger="databricks.sdk.credentials_provider"): + cmd = credentials_provider.DatabricksCliTokenSource._resolve_cli_command(_CLI, cfg) + assert cmd == [_CLI, "auth", "token", "--host", _HOST] + assert any( + "Failed to parse Databricks CLI version" in rec.message and rec.levelname == "DEBUG" for rec in caplog.records + ) + assert not any("does not support --profile" in rec.message for rec in caplog.records) + + # Tests for cloud-agnostic hosts and removed cloud checks class TestCloudAgnosticHosts: """Tests that credential providers work with cloud-agnostic hosts after removing is_azure/is_gcp checks."""