Skip to content

Commit 193da08

Browse files
committed
fix(kernel): guard JWT M2M + databricks-oauth auth_type; renumber comments
Address peco-review-bot review on #921: - Medium: add the missing ambiguity guard for oauth_jwt_key_file + auth_type="databricks-oauth" (U2M intent), mirroring the existing shared-secret M2M + U2M guard. Fails loudly rather than silently resolving to one flow. Covered by a new unit test. - Low: renumber the inline resolution-order comments (PAT→3, U2M→4, creds→5, else→6) to match the docstring after the JWT branch insert. Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
1 parent 1ac6d80 commit 193da08

2 files changed

Lines changed: 25 additions & 4 deletions

File tree

src/databricks/sql/backend/kernel/auth_bridge.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,13 @@ def kernel_auth_kwargs(
248248
"kernel-managed JWT private-key M2M, or use the Thrift backend "
249249
"(default) for credentials_provider."
250250
)
251+
if has_jwt_m2m and auth_type == "databricks-oauth":
252+
raise NotSupportedError(
253+
f"Ambiguous auth on use_kernel=True: auth_type={auth_type!r} selects "
254+
"the U2M browser flow, but oauth_jwt_key_file was also provided "
255+
"(JWT private-key M2M). Drop oauth_jwt_key_file for U2M, or drop "
256+
"auth_type for JWT M2M."
257+
)
251258

252259
# 1. OAuth M2M (JWT private-key client assertion) — the kernel signs a
253260
# short-lived assertion with the private key and runs the
@@ -309,7 +316,7 @@ def kernel_auth_kwargs(
309316
kwargs["identity_federation_client_id"] = federation_client_id
310317
return kwargs
311318

312-
# 2. PAT (including TokenFederationProvider-wrapped PAT).
319+
# 3. PAT (including TokenFederationProvider-wrapped PAT).
313320
if _is_pat(auth_provider):
314321
token = _extract_bearer_token(auth_provider)
315322
if not token:
@@ -322,7 +329,7 @@ def kernel_auth_kwargs(
322329
kwargs["identity_federation_client_id"] = federation_client_id
323330
return kwargs
324331

325-
# 3. OAuth U2M — browser authorization-code flow; the kernel runs it.
332+
# 4. OAuth U2M — browser authorization-code flow; the kernel runs it.
326333
# Only databricks-oauth reaches here (azure-oauth rejected up front).
327334
# Forward the connector's own databricks-sql-python bundle instead of
328335
# the kernel's databricks-sql-connector default, for parity with the
@@ -354,7 +361,7 @@ def kernel_auth_kwargs(
354361
kwargs["identity_federation_client_id"] = federation_client_id
355362
return kwargs
356363

357-
# 4. Custom credentials_provider — the connector's primary M2M path
364+
# 5. Custom credentials_provider — the connector's primary M2M path
358365
# on Thrift/SEA, but unusable on the kernel: it's an opaque token
359366
# source with no extractable client_id/secret, so the kernel
360367
# can't own the token lifecycle. Point the caller at the raw
@@ -368,7 +375,7 @@ def kernel_auth_kwargs(
368375
"credentials_provider."
369376
)
370377

371-
# 5. Everything else (including no usable credentials at all —
378+
# 6. Everything else (including no usable credentials at all —
372379
# ``auth_provider`` is None on the kernel path when no access
373380
# token was supplied and no OAuth kwargs resolved above).
374381
provider_desc = (

tests/unit/test_kernel_auth_bridge.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,20 @@ def test_jwt_plus_credentials_provider_is_rejected(self):
355355
},
356356
)
357357

358+
def test_jwt_plus_databricks_oauth_auth_type_is_rejected(self):
359+
# auth_type="databricks-oauth" signals U2M intent; a private key
360+
# alongside it is ambiguous (mirrors the shared-secret M2M + U2M guard).
361+
with pytest.raises(NotSupportedError, match="oauth_jwt_key_file"):
362+
kernel_auth_kwargs(
363+
None,
364+
{
365+
"oauth_client_id": "sp",
366+
"oauth_jwt_key_file": "/k.pem",
367+
"oauth_jwt_kid": "k",
368+
"auth_type": "databricks-oauth",
369+
},
370+
)
371+
358372
def test_federation_client_id_forwarded(self):
359373
kwargs = kernel_auth_kwargs(
360374
_FakeOAuthProvider(),

0 commit comments

Comments
 (0)