feat(kernel): wire OAuth (M2M/U2M), TLS/mTLS, and Geometry on use_kernel path by vikrantpuppala · Pull Request #819 · databricks/databricks-sql-python

vikrantpuppala · 2026-06-02T10:34:25Z

What

Wires the OAuth + TLS surface (added to the Rust kernel's pyo3 binding in databricks-sql-kernel) through the use_kernel connector path, plus Geometry/Geography result-type mapping. This brings the kernel backend toward parity with Thrift/SEA for auth + TLS.

Companion to the kernel-side PR (databricks-sql-kernel #107). KERNEL_REV is bumped to the kernel commit that ships the new Session kwargs.

Changes

OAuth — kernel-owned token lifecycle

auth_bridge.kernel_auth_kwargs now takes the raw connect() auth options. The OAuth secret is consumed during auth_provider construction and isn't recoverable from the built provider, so the bridge reads it from the original kwargs:
- M2M:oauth_client_id + oauth_client_secret → kernel auth_type='oauth-m2m'. The kernel acquires + refreshes tokens via workspace OIDC client-credentials. oauth_client_secret is a new connect() kwarg, honored only on use_kernel=True (Thrift/SEA unaffected).
- U2M:auth_typedatabricks-oauth / azure-oauth → kernel auth_type='oauth-u2m' (forwards optional oauth_client_id / oauth_redirect_port); the kernel runs the browser flow.
- Custom credentials_provider: rejected with a clear NotSupportedError. It's an opaque token source with no extractable raw creds, so the kernel can't own the lifecycle — callers are pointed at oauth_client_id / oauth_client_secret.
session.py forwards the raw auth kwargs into KernelDatabricksClient.

TLS / mTLS

client._kernel_tls_kwargs translates the connector's SSLOptions → kernel tls_*Session kwargs: CA file → tls_ca_cert (PEM bytes), client cert/key files → tls_client_cert / tls_client_key (mTLS), and inverts tls_verify→tls_skip_verify, tls_verify_hostname→tls_skip_hostname_verify. A password-protected client key is rejected (no kernel surface for it yet). ssl_options is no longer accept-and-ignored on this path.

Geometry

type_mapping recovers GEOGRAPHY / GEOMETRY result columns from the databricks.type_name metadata (they arrive as Arrow Utf8, like VARIANT), so PEP-249 description reports the precise type instead of string.

Design note — why delegate raw creds rather than reuse the connector's OAuth provider

The connector's auth providers are opaque token sources that mint their own tokens; the OAuth client_secret is discarded after provider construction. Rather than run two OAuth stacks (or snapshot a non-refreshing token), we forward the raw credentials to the kernel and let it own acquisition + refresh — matching the kernel-owned model the pyo3 binding was built for. A custom credentials_provider (the connector's M2M path on Thrift/SEA) can't be delegated this way, so it's explicitly rejected on the kernel path.

Tests

140 kernel unit tests pass — new coverage: M2M/U2M routing, M2M-takes-precedence-over-PAT, scope normalization, custom-provider rejection, SSLOptions→tls_* translation (incl. mTLS cert/key, cert-only fallback, encrypted-key rejection, missing-file error), and GEOGRAPHY/GEOMETRY type recovery.
black + mypy clean on the kernel backend.

Caveat for reviewers

KERNEL_REV currently points at the kernel's feat/pyo3-oauth-tls branch HEAD (companion PR #107, not yet merged) so this PR's CI can build against the new Session kwargs. Re-pin to the squash-merge SHA once #107 lands before this merges.

This pull request and its description were written by Isaac.

…nel path Consume the OAuth + TLS surface added to the Rust kernel's pyo3 binding so the use_kernel backend reaches feature parity with the Thrift/SEA backends for these flows. Bumps KERNEL_REV to the kernel commit that ships the new Session kwargs. OAuth (kernel-owned lifecycle): - auth_bridge.kernel_auth_kwargs now takes the raw connect() auth options (the OAuth secret is consumed during auth_provider construction and isn't recoverable from the built provider, so we read it from the original kwargs): - OAuth M2M: oauth_client_id + oauth_client_secret are forwarded to the kernel's auth_type='oauth-m2m'; the kernel acquires and refreshes tokens itself via workspace OIDC client-credentials. - OAuth U2M: auth_type 'databricks-oauth' / 'azure-oauth' route to the kernel's auth_type='oauth-u2m' (optional oauth_client_id / oauth_redirect_port forwarded); the kernel runs the browser flow. - A custom credentials_provider is rejected with a clear NotSupportedError — it's an opaque token source with no extractable raw creds, so the kernel can't own the lifecycle. Callers are pointed at oauth_client_id / oauth_client_secret. - oauth_client_secret is a new connect() kwarg, honored only on the use_kernel path; Thrift/SEA are unaffected. - session.py forwards the raw auth kwargs into KernelDatabricksClient. TLS / mTLS: - client._kernel_tls_kwargs translates the connector's SSLOptions into the kernel's tls_* Session kwargs: tls_trusted_ca_file -> tls_ca_cert (PEM bytes), client cert/key files -> tls_client_cert / tls_client_key (mTLS), and inverts tls_verify -> tls_skip_verify / tls_verify_hostname -> tls_skip_hostname_verify. A password-protected client key is rejected (the kernel has no surface for it yet). ssl_options is no longer accept-and-ignored on this path. Geometry: - type_mapping recovers GEOGRAPHY / GEOMETRY result columns from the databricks.type_name metadata (they arrive as Arrow Utf8, like VARIANT) so PEP-249 description reports the precise type instead of collapsing to string. Tests: 140 kernel unit tests pass (new M2M/U2M routing, custom-provider rejection, SSLOptions->tls_* translation incl. mTLS + encrypted-key rejection, and geospatial type recovery). black + mypy clean. Note: KERNEL_REV points at the kernel's feat/pyo3-oauth-tls branch HEAD until that PR merges; re-pin to the squash-merge SHA before this lands. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

The companion kernel PR (#107) was amended with review fixes (doc backend-inversion fix, U2M degenerate-value validation), changing its branch HEAD SHA. Re-pin so this PR's kernel-e2e builds against the current kernel. Still the branch HEAD — re-pin to the squash-merge SHA once #107 lands. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

gopalldb · 2026-06-02T11:16:53Z

🔴 P0 (Must Fix)

KERNEL_REV pins an unmerged branch HEAD. KERNEL_REV:1 points at 774569bc... which is 1 commit ahead of the kernel's main. CI only validates 40-hex format, not reachability — nothing catches a force-push or
branch deletion that orphans this SHA. Must wait for the companion kernel PR to merge, then re-pin to the squash-merge SHA. Optional follow-up: add git merge-base --is-ancestor check in kernel-e2e CI.

🟡 P1 (Should Fix)

OAuth client_secret retained on long-lived connector (client.py:106, read at :178). Stored as self._auth_options in init, never scrubbed even after open_session() succeeds. The PR description's
"secret is consumed during provider construction" contract is true on Thrift/SEA but false here. Trivially leaks via vars(conn), pickle, debugger inspection. Fix: scrub
self._auth_options.pop("oauth_client_secret", None) in open_session's finally.
tls_verify=False doesn't also imply skip-hostname (client.py:684-687). SSLOptions treats tls_verify=False as a superset of tls_verify_hostname=False; the bridge maps them independently. Depending on
kernel interpretation, "fully disable TLS" intent either fails on hostname or over-disables. Fix: when tls_verify=False, emit both tls_skip_verify=True AND tls_skip_hostname_verify=True.
credentials_provider + M2M kwargs silently picks M2M (auth_bridge.py:147-159). The rejection at :190 is unreachable when M2M kwargs are also present — M2M wins at step 1. PR description called this out as
a focus area; implementation doesn't enforce it. Fix: detect ambiguity before resolving.
auth_type='databricks-oauth' + oauth_client_secret silently uses M2M (same routing bug, different combo). User asked for U2M, gets M2M against the wrong principal — failure surfaces only as a confusing
first-call 401.
No secret-leak tests. Brand-new credential kwarg, zero tests asserting it's absent from caplog, exception str/repr, or vars(connector). _errors.py:wrap_kernel_exception:133 formats {exc!r} into an
OperationalError — any kernel-side error embedding arg values would leak through and CI wouldn't notice.

🟢 P2 (Nice-to-have)

Empty/unreadable CA file passes through (cryptic kernel error)
_normalize_scopes returns None for unknown types instead of raising
Geometry/Geography mapping depends on undocumented databricks.type_name stability
U2M oauth_scopes is forwarded but untested
auth_kwargs/tls_kwargs computed before the try-block — secret materialised in a local that's outside the scrub finally

@gopalldb

… ambiguity, TLS superset Addresses @gopalldb's review on #819. P0: - Re-pin KERNEL_REV to the squash-merge SHA of the now-merged kernel PR (ec2288742cbac0cd9fab50da353e8405972eefe9 on kernel main), replacing the orphaned branch-HEAD SHA. P1: - Scrub oauth_client_secret from the long-lived self._auth_options in open_session's finally (even on the failure path). It outlives the method, so a retained secret was exposed to vars()/pickle/debugger far longer than needed; the kernel now owns it. - tls_verify=False now also emits tls_skip_hostname_verify=True — no-verify subsumes hostname verification, matching SSLOptions' create_ssl_context (check_hostname=False when tls_verify is False). - Reject ambiguous auth combos before resolving, instead of silently picking M2M: (a) credentials_provider + oauth_client_secret, and (b) a U2M auth_type (databricks-oauth/azure-oauth) + oauth_client_secret. Both raise NotSupportedError naming the conflict, so the failure is at session-open rather than a confusing first-call 401 against the wrong principal. - Secret-leak tests: oauth_client_secret is forwarded to the kernel but scrubbed from _auth_options, and absent from repr(conn) / vars(conn); scrub runs even when open_session raises. P2: - _read_pem_bytes rejects an empty/whitespace CA/cert file with a clear ProgrammingError instead of passing empty PEM to the kernel. - _normalize_scopes raises ProgrammingError on a non-str/list/tuple oauth_scopes rather than silently dropping to default scopes. - Added U2M oauth_scopes forwarding test. 150 kernel unit tests pass; black + mypy clean. Note: KERNEL_REV now points at the merged kernel main SHA; no further re-pin needed before merge. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

vikrantpuppala · 2026-06-02T11:52:25Z

Thanks @gopalldb — all P0/P1 addressed plus the actionable P2s, in eed76ff6:

P0 — KERNEL_REV

Re-pinned to the squash-merge SHA ec2288742cbac0cd9fab50da353e8405972eefe9 (now on kernel main), replacing the orphaned branch-HEAD SHA. No further re-pin needed before merge. (Good call on the git merge-base --is-ancestor CI check — filing that as a kernel-e2e follow-up.)

P1

Secret retention — open_session's finally now scrubs oauth_client_secret from the long-lived self._auth_options (runs even on the failure path). You're right the "consumed during provider construction" contract didn't hold here.
tls_verify=False ⊃ hostname — now emits both tls_skip_verifyandtls_skip_hostname_verify, matching SSLOptions.create_ssl_context (which sets check_hostname=False whenever tls_verify is False).
credentials_provider + M2M — now rejected up front with a clear NotSupportedError (was unreachable behind the step-1 M2M win).
databricks-oauth + oauth_client_secret — same ambiguity class, now rejected naming the conflict instead of silently routing M2M against the wrong principal.
Secret-leak tests — added: secret forwarded to the kernel but scrubbed from _auth_options, and absent from repr(conn) / vars(conn); scrub verified on the raise path too.

P2

Empty/whitespace CA/cert file → clear ProgrammingError instead of cryptic kernel parse error.
_normalize_scopes raises on non-str/list/tuple instead of silently defaulting.
Added a U2M oauth_scopes forwarding test.
Left as-is with rationale: the databricks.type_name geometry dependency mirrors the existing VARIANT/complex handling (same metadata contract), and the pre-try materialization is covered by the new _auth_options scrub + local-dict scrub.

150 kernel unit tests pass; black + mypy clean.

…+ add TLS e2e Live mitmproxy TLS e2e (this commit) surfaced a real bug: on use_kernel=True the connector still built its own auth provider in Session.__init__ via get_python_sql_connector_auth_provider, BEFORE use_kernel was consulted. For OAuth that constructor eagerly runs the flow (DatabricksOAuthProvider calls _initial_get_token in __init__), so passing oauth_client_id for kernel-managed M2M fired the U2M *browser* flow at connect() time — opening a browser and using the M2M service principal id as a U2M OAuth app client_id (which the account rejects). Fix: on use_kernel=True, do NOT build the connector's provider. Hand the kernel auth bridge a minimal AccessTokenAuthProvider when an access_token is present, else None; OAuth M2M/U2M resolve purely from the raw connect() kwargs the bridge already reads, and the kernel owns the entire token lifecycle. Thrift/SEA backends are unchanged. - session.py: branch the provider build on use_kernel; import AccessTokenAuthProvider; update the kernel-branch comment. - auth_bridge.py: kernel_auth_kwargs / _is_pat / _extract_bearer_token accept Optional[AuthProvider] (None when no token on the kernel path); clearer "no credentials" error; refreshed module docstring. - tests/unit/test_session.py: regression guards — use_kernel must NOT call get_python_sql_connector_auth_provider, forwards raw M2M creds via auth_options, and uses a minimal AccessTokenAuthProvider for PAT. TLS e2e (the connector-side counterpart to the kernel's v0_tls_e2e.rs): - tests/e2e/test_kernel_tls.py: 3 mitmproxy-backed tests driving sql.connect(use_kernel=True, _tls_*=...) — strict default rejects the re-signed cert, _tls_trusted_ca_file trusts it, _tls_no_verify bypasses. Skips cleanly without the kernel wheel / creds / MITMPROXY_CA_CERT. Verified live against the dogfood workspace on the OAuth M2M path (all 3 green) — this is what caught the bug above. - kernel-e2e.yml: run test_kernel_tls.py in the existing (label/merge- queue-gated) run-kernel-e2e job behind a mitmproxy service. The kernel wheel is built before the proxy is in scope so cargo's JFrog fetch isn't routed through it; HTTPS_PROXY is scoped to the TLS test step. Path filter now also triggers on session.py and test_kernel_tls.py. 184 unit tests pass; black + mypy clean. Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

The two TestKernelAuthProviderBypass tests went through databricks.sql.connect(use_kernel=True), which triggers the lazy `import databricks_sql_kernel` in the kernel client module. That import fails in the unit-test job (no Rust wheel installed), so the tests errored with ImportError before their patches applied — green locally (wheel present) but red in CI. Rewrite them to exercise Session.__init__'s provider-selection logic directly with _create_backend stubbed, so the kernel client (and its import) is never touched. Assert on session.auth_provider and that get_python_sql_connector_auth_provider is not called. Verified passing both with the wheel present and with `import databricks_sql_kernel` forcibly blocked (reproducing the CI environment). Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

@gopalldb

…aders Addresses @gopalldb's review on #823. P1.1 — secret scrub now covers a kwarg-build failure: - open_session built auth/tls/retry/http_headers kwargs OUTSIDE the try whose finally scrubs oauth_client_secret. A raise in kernel_auth_kwargs (e.g. OAuth token-exchange failure with the M2M secret in hand) would skip the scrub, leaking the secret onto the long-lived connector. Moved all kwarg builds inside the try (auth_kwargs / tls_kwargs pre-declared empty so the finally can reference them on an early raise). P1.3 / P2 — don't forward kernel-managed headers: - The connector now drops Authorization and x-databricks-org-id from http_headers before forwarding to the kernel (new _KERNEL_MANAGED_HEADERS set). The kernel manages both (auth from the provider; org-id re-derived from ?o= in http_path), so forwarding is redundant — and the connector always injects the SPOG org-id, which the kernel skips-and-warns per request, so this also kills the WARN spam. Double-walls the kernel's own reserved-name skip. Tests: - unit: Authorization / x-databricks-org-id dropped before forwarding; only-reserved-headers omits the kwarg entirely (test_kernel_client). - P1.2: user_agent_entry reaches the kernel client's http_headers on use_kernel=True (test_session) — guards a regression where session.py stops folding the entry into the composed User-Agent. CHANGELOG: deferred to release-cut per repo convention (entries are added at version bump, not per-PR; #819/#820 followed the same). 197 kernel unit tests pass; black + mypy clean; 3 e2e pass live. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

@gopalldb

… use_kernel (databricks#823) * feat(kernel): wire statement query tags + http_headers/user_agent on use_kernel Consume the kernel surface from kernel PR (query tags + custom HTTP headers) on the use_kernel path. Bumps KERNEL_REV. Query tags (statement-level): - execute_command no longer raises NotSupportedError for query_tags; it calls stmt.set_query_tags(query_tags) after set_sql. The connector already passes Dict[str, Optional[str]], which the kernel accepts (None value -> bare key in the SEA query_tags conf). http_headers + user_agent_entry: - The kernel client now forwards http_headers to the kernel Session as the `http_headers` kwarg (was accept-and-ignore). session.py already passes all_headers, which carries the connector's composed User-Agent (PyDatabricksSqlConnector/x (entry)) + caller headers + SPOG org-id. - The kernel applies them per request: its own Authorization / org-id win; a caller User-Agent is APPENDED to the kernel base UA (the base carries the DatabricksJDBCDriverOSS token that gates the SEA result disposition, so it's never replaced). So user_agent_entry is honored end-to-end via the existing http_headers forwarding — no separate kwarg needed. Tests: - unit: query_tags forwarded to set_query_tags (was: rejection test); http_headers forwarded to the kernel Session (and omitted when empty). - e2e (test_kernel_backend.py): a query_tagged query and a connection with user_agent_entry + a custom http_header both round-trip green against a dogfood warehouse. The UA case specifically guards the append behavior (replacing the base UA would 400 on the result disposition). KERNEL_REV -> c2053f68b75fef4a29425096dc6bbafb774d8b83 (kernel PR databricks#119 branch HEAD; re-pin to the squash-merge SHA once databricks#119 lands). 194 kernel unit tests pass; black + mypy clean; 3 e2e pass live. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com> * fix(kernel): address review — secret-scrub ordering, drop reserved headers Addresses @gopalldb's review on databricks#823. P1.1 — secret scrub now covers a kwarg-build failure: - open_session built auth/tls/retry/http_headers kwargs OUTSIDE the try whose finally scrubs oauth_client_secret. A raise in kernel_auth_kwargs (e.g. OAuth token-exchange failure with the M2M secret in hand) would skip the scrub, leaking the secret onto the long-lived connector. Moved all kwarg builds inside the try (auth_kwargs / tls_kwargs pre-declared empty so the finally can reference them on an early raise). P1.3 / P2 — don't forward kernel-managed headers: - The connector now drops Authorization and x-databricks-org-id from http_headers before forwarding to the kernel (new _KERNEL_MANAGED_HEADERS set). The kernel manages both (auth from the provider; org-id re-derived from ?o= in http_path), so forwarding is redundant — and the connector always injects the SPOG org-id, which the kernel skips-and-warns per request, so this also kills the WARN spam. Double-walls the kernel's own reserved-name skip. Tests: - unit: Authorization / x-databricks-org-id dropped before forwarding; only-reserved-headers omits the kwarg entirely (test_kernel_client). - P1.2: user_agent_entry reaches the kernel client's http_headers on use_kernel=True (test_session) — guards a regression where session.py stops folding the entry into the composed User-Agent. CHANGELOG: deferred to release-cut per repo convention (entries are added at version bump, not per-PR; databricks#819/databricks#820 followed the same). 197 kernel unit tests pass; black + mypy clean; 3 e2e pass live. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com> * chore: re-pin KERNEL_REV to merged kernel main (101aa46) databricks#119 (statement query tags + custom HTTP headers) is now merged to kernel main. Re-pin from the orphaned branch HEAD (c2053f6) to the current merged main HEAD (101aa46), which contains databricks#119's merge commit (df8302f) and is reachable from main — no orphan-SHA risk. Verified end-to-end against a wheel built from 101aa46: kernel e2e (query_tags_round_trip, user_agent_entry + http_headers round_trip, select_one) all pass. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com> --------- Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 10:34 — with GitHub Actions Inactive

vikrantpuppala added the kernel-e2e Trigger preview run of the Kernel E2E workflow on this PR label Jun 2, 2026

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 10:41 — with GitHub Actions Inactive

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 10:59 — with GitHub Actions Inactive

github-actionsBot removed the kernel-e2e Trigger preview run of the Kernel E2E workflow on this PR label Jun 2, 2026

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 11:52 — with GitHub Actions Inactive

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 12:21 — with GitHub Actions Inactive

vikrantpuppala added the kernel-e2e Trigger preview run of the Kernel E2E workflow on this PR label Jun 2, 2026

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 12:26 — with GitHub Actions Inactive

vikrantpuppala temporarily deployed to azure-prod June 2, 2026 12:42 — with GitHub Actions Inactive

github-actionsBot removed the kernel-e2e Trigger preview run of the Kernel E2E workflow on this PR label Jun 2, 2026

gopalldb approved these changes Jun 2, 2026
View reviewed changes

vikrantpuppala added this pull request to the merge queue Jun 2, 2026

Merged via the queue into main with commit 39dc0a7Jun 2, 2026
46 checks passed

msrathore-db mentioned this pull request Jun 2, 2026
Bump to version 4.2.7 #821
Merged
3 tasks

gopalldb mentioned this pull request Jun 3, 2026
feat(kernel): statement query tags + http_headers/user_agent_entry on use_kernel #823
Merged

Navigation Menu

Uh oh!

feat(kernel): wire OAuth (M2M/U2M), TLS/mTLS, and Geometry on use_kernel path - #819

feat(kernel): wire OAuth (M2M/U2M), TLS/mTLS, and Geometry on use_kernel path#819
vikrantpuppala merged 5 commits into
mainfrom
feat/use-kernel-oauth-tls-geometry

vikrantpuppala commented Jun 2, 2026

Uh oh!

gopalldb commented Jun 2, 2026

Uh oh!

vikrantpuppala commented Jun 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

vikrantpuppala commented Jun 2, 2026

What

Changes

Design note — why delegate raw creds rather than reuse the connector's OAuth provider

Tests

Caveat for reviewers

Uh oh!

gopalldb commented Jun 2, 2026

Uh oh!

vikrantpuppala commented Jun 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants