Skip to content

Commit 61c4bdc

Browse files
Add public helper to build OIDC credentials providers. (#1036)
## What changes are proposed in this pull request? This PR makes the `oidc_credentials_provider` function public to make it easier for users to define their own credentials provider from a custom OIDC token source. ## How is this tested? Unit tests.
1 parent 444e27b commit 61c4bdc

File tree

3 files changed

+10
-7
lines changed

3 files changed

+10
-7
lines changed

NEXT_CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
### New Features and Improvements
66

7+
* Add a public helper function to build a `CredentialsProvider` directly from an `IdTokenSource`.
8+
79
### Bug Fixes
810

911
### Documentation

databricks/sdk/credentials_provider.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -319,17 +319,18 @@ def env_oidc(cfg) -> Optional[CredentialsProvider]:
319319
if cfg.oidc_token_env:
320320
env_var = cfg.oidc_token_env
321321

322-
return _oidc_credentials_provider(cfg, oidc.EnvIdTokenSource(env_var))
322+
return oidc_credentials_provider(cfg, oidc.EnvIdTokenSource(env_var))
323323

324324

325325
@credentials_strategy("file-oidc", ["host", "oidc_token_filepath"])
326326
def file_oidc(cfg) -> Optional[CredentialsProvider]:
327-
return _oidc_credentials_provider(cfg, oidc.FileIdTokenSource(cfg.oidc_token_filepath))
327+
return oidc_credentials_provider(cfg, oidc.FileIdTokenSource(cfg.oidc_token_filepath))
328328

329329

330-
# This function is a helper function to create an OIDC CredentialsProvider
331-
# that provides a Databricks token from an IdTokenSource.
332-
def _oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optional[CredentialsProvider]:
330+
def oidc_credentials_provider(cfg, id_token_source: oidc.IdTokenSource) -> Optional[CredentialsProvider]:
331+
"""Creates a CredentialsProvider to sign requests with an OAuth token obtained
332+
by automatically performing the token exchange using the given IdTokenSource."""
333+
333334
try:
334335
id_token_source.id_token() # validate the id_token_source
335336
except Exception as e:

tests/test_credentials_provider.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ def test_oidc_credentials_provider_invalid_id_token_source():
188188
id_token_source = Mock()
189189
id_token_source.id_token.side_effect = ValueError("Invalid ID token source")
190190

191-
cp = credentials_provider._oidc_credentials_provider(mock_cfg, id_token_source)
191+
cp = credentials_provider.oidc_credentials_provider(mock_cfg, id_token_source)
192192
assert cp is None
193193

194194

@@ -216,7 +216,7 @@ def mock_exchange_id_token(id_token: oidc.IdToken):
216216

217217
mocker.patch.object(oidc.DatabricksOidcTokenSource, "_exchange_id_token", side_effect=mock_exchange_id_token)
218218

219-
cp = credentials_provider._oidc_credentials_provider(mock_cfg, id_token_source)
219+
cp = credentials_provider.oidc_credentials_provider(mock_cfg, id_token_source)
220220
assert cp is not None
221221

222222
# Test that the credentials provider returns the expected headers

0 commit comments

Comments
 (0)