Skip to content

Commit 4fb6bd0

Browse files
config changes
1 parent 5cc74ef commit 4fb6bd0

File tree

3 files changed

+93
-1
lines changed

3 files changed

+93
-1
lines changed

databricks/sdk/config.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import logging
55
import os
66
import pathlib
7+
import re
78
import sys
89
import urllib.parse
910
from enum import Enum
@@ -137,6 +138,10 @@ class Config:
137138
scopes: str = ConfigAttribute()
138139
authorization_details: str = ConfigAttribute()
139140

141+
# Controls whether the offline_access scope is requested during U2M OAuth authentication.
142+
# offline_access is requested by default, causing a refresh token to be included in the OAuth token.
143+
disable_oauth_refresh_token: bool = ConfigAttribute(env="DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN")
144+
140145
files_ext_client_download_streaming_chunk_size: int = 2 * 1024 * 1024 # 2 MiB
141146

142147
# When downloading a file, the maximum number of attempts to retry downloading the whole file. Default is no limit.
@@ -265,6 +270,7 @@ def __init__(
265270
self._known_file_config_loader()
266271
self._fix_host_if_needed()
267272
self._validate()
273+
self._sort_scopes()
268274
self.init_auth()
269275
self._init_product(product, product_version)
270276
except ValueError as e:
@@ -666,6 +672,16 @@ def _validate(self):
666672
names = " and ".join(sorted(auths_used))
667673
raise ValueError(f"validate: more than one authorization method configured: {names}")
668674

675+
def _sort_scopes(self):
676+
"""Sort scopes in-place for better de-duplication in the refresh token cache.
677+
Delimiter is set to a single whitespace after sorting."""
678+
if self.scopes and isinstance(self.scopes, str):
679+
# Split on whitespaces and commas, sort, and rejoin
680+
parsed = [s for s in re.split(r"[\s,]+", self.scopes) if s]
681+
if parsed:
682+
parsed.sort()
683+
self.scopes = " ".join(parsed)
684+
669685
def init_auth(self):
670686
try:
671687
self._header_factory = self._credentials_strategy(self)
@@ -685,6 +701,33 @@ def _init_product(self, product, product_version):
685701
else:
686702
self._product_info = None
687703

704+
def get_scopes(self) -> List[str]:
705+
"""Get OAuth scopes with proper defaulting.
706+
707+
Returns ["all-apis"] if no scopes configured.
708+
This is the single source of truth for scope defaulting across all OAuth methods.
709+
710+
Parses string scopes by splitting on whitespaces and commas.
711+
712+
Returns:
713+
List of scope strings.
714+
"""
715+
if self.scopes and isinstance(self.scopes, str):
716+
parsed = [s for s in re.split(r"[\s,]+", self.scopes) if s]
717+
if not parsed: # Empty string case
718+
return ["all-apis"]
719+
return parsed
720+
return ["all-apis"]
721+
722+
def get_scopes_as_string(self) -> str:
723+
"""Get OAuth scopes as a space-separated string.
724+
725+
Returns "all-apis" if no scopes configured.
726+
"""
727+
if self.scopes and isinstance(self.scopes, str):
728+
return self.scopes
729+
return " ".join(self.get_scopes())
730+
688731
def __repr__(self):
689732
return f"<{self.debug_string()}>"
690733

tests/test_config.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -453,3 +453,41 @@ def test_no_org_id_header_on_regular_workspace(requests_mock):
453453

454454
# Verify the X-Databricks-Org-Id header was NOT added
455455
assert "X-Databricks-Org-Id" not in requests_mock.last_request.headers
456+
457+
def test_disable_oauth_refresh_token_from_env(monkeypatch, mocker):
458+
mocker.patch("databricks.sdk.config.Config.init_auth")
459+
monkeypatch.setenv("DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN", "true")
460+
config = Config(host="https://test.databricks.com")
461+
assert config.disable_oauth_refresh_token is True
462+
463+
464+
def test_disable_oauth_refresh_token_defaults_to_false(mocker):
465+
mocker.patch("databricks.sdk.config.Config.init_auth")
466+
config = Config(host="https://test.databricks.com")
467+
assert config.disable_oauth_refresh_token is None # ConfigAttribute returns None when not set
468+
469+
470+
def test_config_file_scopes_empty_defaults_to_all_apis(monkeypatch, mocker):
471+
"""Test that empty scopes in config file defaults to all-apis."""
472+
mocker.patch("databricks.sdk.config.Config.init_auth")
473+
monkeypatch.setenv("HOME", str(pathlib.Path(__tests__) / "testdata"))
474+
config = Config(profile="scope-empty")
475+
assert config.get_scopes() == ["all-apis"]
476+
477+
478+
def test_config_file_scopes_single(monkeypatch, mocker):
479+
"""Test single scope from config file."""
480+
mocker.patch("databricks.sdk.config.Config.init_auth")
481+
monkeypatch.setenv("HOME", str(pathlib.Path(__tests__) / "testdata"))
482+
config = Config(profile="scope-single")
483+
assert config.get_scopes() == ["clusters:read"]
484+
485+
486+
def test_config_file_scopes_multiple_sorted(monkeypatch, mocker):
487+
"""Test multiple scopes from config file are sorted."""
488+
mocker.patch("databricks.sdk.config.Config.init_auth")
489+
monkeypatch.setenv("HOME", str(pathlib.Path(__tests__) / "testdata"))
490+
config = Config(profile="scope-multiple")
491+
# Should be sorted alphabetically
492+
expected = ["clusters", "files:read", "iam:read", "jobs", "mlflow", "model-serving:read", "pipelines"]
493+
assert config.get_scopes() == expected

tests/testdata/.databrickscfg

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,4 +38,15 @@ google_credentials = paw48590aw8e09t8apu
3838

3939
[pat.with.dot]
4040
host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/
41-
token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ
41+
token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ
42+
43+
[scope-empty]
44+
host = https://example.cloud.databricks.com
45+
46+
[scope-single]
47+
host = https://example.cloud.databricks.com
48+
scopes = clusters
49+
50+
[scope-multiple]
51+
host = https://example.cloud.databricks.com
52+
scopes = clusters, jobs, pipelines, iam:read, files:read, mlflow, model-serving:read

0 commit comments

Comments
 (0)