Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions databricks/sdk/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import logging
import os
import pathlib
import re
import sys
import urllib.parse
from enum import Enum
Expand Down Expand Up @@ -137,6 +138,10 @@ class Config:
scopes: str = ConfigAttribute()
authorization_details: str = ConfigAttribute()

# Controls whether the offline_access scope is requested during U2M OAuth authentication.
# offline_access is requested by default, causing a refresh token to be included in the OAuth token.
disable_oauth_refresh_token: bool = ConfigAttribute(env="DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN")

files_ext_client_download_streaming_chunk_size: int = 2 * 1024 * 1024 # 2 MiB

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

def _sort_scopes(self):
"""Sort scopes in-place for better de-duplication in the refresh token cache.
Delimiter is set to a single whitespace after sorting."""
if self.scopes and isinstance(self.scopes, str):
# Split on whitespaces and commas, sort, and rejoin
parsed = [s for s in re.split(r"[\s,]+", self.scopes) if s]
if parsed:
parsed.sort()
self.scopes = " ".join(parsed)

def init_auth(self):
try:
self._header_factory = self._credentials_strategy(self)
Expand All @@ -685,6 +701,33 @@ def _init_product(self, product, product_version):
else:
self._product_info = None

def get_scopes(self) -> List[str]:
"""Get OAuth scopes with proper defaulting.

Returns ["all-apis"] if no scopes configured.
This is the single source of truth for scope defaulting across all OAuth methods.

Parses string scopes by splitting on whitespaces and commas.

Returns:
List of scope strings.
"""
if self.scopes and isinstance(self.scopes, str):
parsed = [s for s in re.split(r"[\s,]+", self.scopes) if s]
if not parsed: # Empty string case
return ["all-apis"]
return parsed
return ["all-apis"]

def get_scopes_as_string(self) -> str:
"""Get OAuth scopes as a space-separated string.

Returns "all-apis" if no scopes configured.
"""
if self.scopes and isinstance(self.scopes, str):
return self.scopes
return " ".join(self.get_scopes())

def __repr__(self):
return f"<{self.debug_string()}>"

Expand Down
41 changes: 40 additions & 1 deletion tests/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
with_user_agent_extra)
from databricks.sdk.version import __version__

from .conftest import noop_credentials, set_az_path
from .conftest import noop_credentials, set_az_path, set_home

__tests__ = os.path.dirname(__file__)

Expand Down Expand Up @@ -453,3 +453,42 @@ def test_no_org_id_header_on_regular_workspace(requests_mock):

# Verify the X-Databricks-Org-Id header was NOT added
assert "X-Databricks-Org-Id" not in requests_mock.last_request.headers


def test_disable_oauth_refresh_token_from_env(monkeypatch, mocker):
mocker.patch("databricks.sdk.config.Config.init_auth")
monkeypatch.setenv("DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN", "true")
config = Config(host="https://test.databricks.com")
assert config.disable_oauth_refresh_token is True


def test_disable_oauth_refresh_token_defaults_to_false(mocker):
mocker.patch("databricks.sdk.config.Config.init_auth")
config = Config(host="https://test.databricks.com")
assert config.disable_oauth_refresh_token is None # ConfigAttribute returns None when not set


def test_config_file_scopes_empty_defaults_to_all_apis(monkeypatch, mocker):
"""Test that empty scopes in config file defaults to all-apis."""
mocker.patch("databricks.sdk.config.Config.init_auth")
set_home(monkeypatch, "/testdata")
config = Config(profile="scope-empty")
assert config.get_scopes() == ["all-apis"]


def test_config_file_scopes_single(monkeypatch, mocker):
"""Test single scope from config file."""
mocker.patch("databricks.sdk.config.Config.init_auth")
set_home(monkeypatch, "/testdata")
config = Config(profile="scope-single")
assert config.get_scopes() == ["clusters"]


def test_config_file_scopes_multiple_sorted(monkeypatch, mocker):
"""Test multiple scopes from config file are sorted."""
mocker.patch("databricks.sdk.config.Config.init_auth")
set_home(monkeypatch, "/testdata")
config = Config(profile="scope-multiple")
# Should be sorted alphabetically
expected = ["clusters", "files:read", "iam:read", "jobs", "mlflow", "model-serving:read", "pipelines"]
assert config.get_scopes() == expected
2 changes: 1 addition & 1 deletion tests/test_notebook_oauth.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def test_config_authenticate_integration(

@pytest.mark.parametrize(
"scopes_input,expected_scopes",
[(["sql", "offline_access"], "sql offline_access")],
[(["sql", "offline_access"], "offline_access sql")],
)
def test_workspace_client_integration(
mock_runtime_env, mock_runtime_native_auth, mock_pat_exchange, scopes_input, expected_scopes
Expand Down
13 changes: 12 additions & 1 deletion tests/testdata/.databrickscfg
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,15 @@ google_credentials = paw48590aw8e09t8apu

[pat.with.dot]
host = https://dbc-XXXXXXXX-YYYY.cloud.databricks.com/
token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ
token = PT0+IC9kZXYvdXJhbmRvbSA8PT0KYFZ

[scope-empty]
host = https://example.cloud.databricks.com

[scope-single]
host = https://example.cloud.databricks.com
scopes = clusters

[scope-multiple]
host = https://example.cloud.databricks.com
scopes = clusters, jobs, pipelines, iam:read, files:read, mlflow, model-serving:read
Loading