fix(kernel): don't build connector OAuth provider on use_kernel path + 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>

Author: vikrantpuppala

.github/workflows/kernel-e2e.yml

@@ -178,7 +178,7 @@ jobs:
           echo "$CHANGED"
           # Run when the connector kernel backend, kernel e2e tests,
           # this workflow, the kernel revision pin, or core deps move.
-          if echo "$CHANGED" | grep -qE "^(src/databricks/sql/backend/kernel/|tests/e2e/test_kernel_backend\.py|tests/unit/test_kernel_|\.github/workflows/kernel-e2e\.yml|KERNEL_REV|pyproject\.toml|poetry\.lock)"; then
+          if echo "$CHANGED" | grep -qE "^(src/databricks/sql/backend/kernel/|src/databricks/sql/session\.py|tests/e2e/test_kernel_backend\.py|tests/e2e/test_kernel_tls\.py|tests/unit/test_kernel_|\.github/workflows/kernel-e2e\.yml|KERNEL_REV|pyproject\.toml|poetry\.lock)"; then
             echo "run_tests=true" >> "$GITHUB_OUTPUT"
           else
             echo "run_tests=false" >> "$GITHUB_OUTPUT"
@@ -319,6 +319,58 @@ jobs:
       - name: Run kernel e2e tests
         run: poetry run pytest tests/e2e/test_kernel_backend.py -v
 
+      # ── TLS e2e through an intercepting proxy ───────────────────────
+      # mitmproxy re-signs every TLS connection with its own CA. The
+      # connector tests (tests/e2e/test_kernel_tls.py) prove the whole
+      # stack — sql.connect(use_kernel=True, _tls_*=...) → SSLOptions →
+      # kernel handshake — trusts the CA / accepts no-verify / rejects
+      # by default. The kernel wheel is already built above, BEFORE any
+      # proxy is in scope, so cargo's JFrog registry fetch was never
+      # routed through (and broken by) mitmproxy.
+      - name: Start mitmproxy
+        run: |
+          docker run -d --name mitmproxy \
+            -p 8080:8080 \
+            mitmproxy/mitmproxy:12.2.1@sha256:743b6cdc817211d64bc269f5defacca8d14e76e647fc474e5c7244dbcb645141 \
+            mitmdump --set stream_large_bodies=0
+          for i in $(seq 1 15); do
+            if docker exec mitmproxy test -f /home/mitmproxy/.mitmproxy/mitmproxy-ca-cert.pem 2>/dev/null; then
+              echo "mitmproxy CA cert is ready"
+              break
+            fi
+            echo "Waiting for mitmproxy CA cert... ($i/15)"
+            sleep 1
+          done
+
+      - name: Extract mitmproxy CA certificate
+        run: |
+          docker cp mitmproxy:/home/mitmproxy/.mitmproxy/mitmproxy-ca-cert.pem /tmp/mitmproxy-ca-cert.pem
+          echo "MITMPROXY_CA_CERT=/tmp/mitmproxy-ca-cert.pem" >> "$GITHUB_ENV"
+          echo "Extracted mitmproxy CA cert:"
+          openssl x509 -in /tmp/mitmproxy-ca-cert.pem -noout -subject -issuer
+
+      - name: Run kernel TLS e2e tests
+        # HTTPS_PROXY is scoped to THIS step so only the test's own HTTP
+        # traffic flows through mitmproxy — not cargo / pip / anything
+        # the earlier build steps needed from the JFrog registry.
+        env:
+          HTTPS_PROXY: http://localhost:8080
+        run: poetry run pytest tests/e2e/test_kernel_tls.py -v
+
+      - name: Verify traffic flowed through mitmproxy
+        if: always()
+        run: |
+          echo "=== mitmproxy logs (tail) ==="
+          docker logs mitmproxy 2>&1 | tail -50
+          echo "=== checking for proxied-connection entries ==="
+          docker logs mitmproxy 2>&1 \
+            | grep -i "CONNECT\|clientconnect\|<<" \
+            || (echo "FAIL: no traffic in mitmproxy logs — proxy was not used" && exit 1)
+
+      - name: Cleanup mitmproxy
+        if: always()
+        run: docker rm -f mitmproxy 2>/dev/null || true
+
       # Post a Kernel E2E check on both the labeled-PR and merge-queue
       # paths so the named check on the PR reflects the latest real
       # run (overwriting the synthetic-success check that
@@ -341,7 +393,7 @@ jobs:
               completed_at: new Date().toISOString(),
               output: {
                 title: 'Kernel E2E passed',
-                summary: 'tests/e2e/test_kernel_backend.py ran green against the pinned kernel SHA.'
+                summary: 'test_kernel_backend.py and the mitmproxy-backed test_kernel_tls.py ran green against the pinned kernel SHA.'
               }
             });
 

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

@@ -30,9 +30,12 @@
 with a clear message rather than deep inside the kernel.
 
 The M2M / U2M decisions are driven by the *raw* connect() kwargs
-(``auth_options``), not the built ``AuthProvider`` — the secret is
-consumed and discarded during provider construction, so it can only be
-read from the original kwargs.
+(``auth_options``), not a built ``AuthProvider``. On the kernel path
+the connector deliberately does **not** build its own OAuth provider
+(that would eagerly run the U2M browser flow / M2M token exchange at
+connect() time, before the kernel is consulted), so ``auth_provider``
+is either a minimal PAT provider or ``None`` and the OAuth credentials
+are available only from the raw kwargs.
 """
 
 from __future__ import annotations
@@ -65,7 +68,7 @@
 _TOKEN_REJECT_RE = re.compile(r"[\x00-\x20\x7f]")
 
 
-def _is_pat(auth_provider: AuthProvider) -> bool:
+def _is_pat(auth_provider: Optional[AuthProvider]) -> bool:
     """Return True iff this provider ultimately wraps an
     ``AccessTokenAuthProvider``.
 
@@ -84,18 +87,20 @@ def _is_pat(auth_provider: AuthProvider) -> bool:
     return False
 
 
-def _extract_bearer_token(auth_provider: AuthProvider) -> Optional[str]:
+def _extract_bearer_token(auth_provider: Optional[AuthProvider]) -> Optional[str]:
     """Pull the current bearer token out of an ``AuthProvider``.
 
     The connector's ``AuthProvider.add_headers`` mutates a header
     dict and writes the ``Authorization: Bearer <token>`` value.
     Going through that public surface keeps us insulated from
     provider-specific internals.
 
-    Returns ``None`` if the provider did not write an Authorization
-    header or wrote a non-Bearer scheme — neither is representable
-    in the kernel's PAT auth surface.
+    Returns ``None`` if there is no provider, the provider did not
+    write an Authorization header, or it wrote a non-Bearer scheme —
+    none of which is representable in the kernel's PAT auth surface.
     """
+    if auth_provider is None:
+        return None
     headers: Dict[str, str] = {}
     auth_provider.add_headers(headers)
     auth = headers.get("Authorization")
@@ -113,7 +118,7 @@ def _extract_bearer_token(auth_provider: AuthProvider) -> Optional[str]:
 
 
 def kernel_auth_kwargs(
-    auth_provider: AuthProvider,
+    auth_provider: Optional[AuthProvider],
     auth_options: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, Any]:
     """Build the kwargs passed to ``databricks_sql_kernel.Session(...)``.
@@ -225,13 +230,18 @@ def kernel_auth_kwargs(
             "credentials_provider."
         )
 
-    # 5. Everything else.
+    # 5. Everything else (including no usable credentials at all —
+    #    ``auth_provider`` is None on the kernel path when no access
+    #    token was supplied and no OAuth kwargs resolved above).
+    provider_desc = (
+        type(auth_provider).__name__ if auth_provider is not None else "no credentials"
+    )
     raise NotSupportedError(
-        f"use_kernel=True supports PAT, OAuth M2M (oauth_client_id + "
-        f"oauth_client_secret), and OAuth U2M (auth_type='databricks-oauth' / "
-        f"'azure-oauth'), but got {type(auth_provider).__name__} with "
-        f"auth_type={auth_type!r}. Use the Thrift backend (default) for other "
-        "auth flows."
+        f"use_kernel=True requires PAT (access_token), OAuth M2M "
+        f"(oauth_client_id + oauth_client_secret), or OAuth U2M "
+        f"(auth_type='databricks-oauth' / 'azure-oauth'), but got "
+        f"{provider_desc} with auth_type={auth_type!r}. Use the Thrift "
+        "backend (default) for other auth flows."
     )
 
 

src/databricks/sql/session.py

@@ -5,6 +5,7 @@
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
+from databricks.sql.auth.authenticators import AccessTokenAuthProvider
 from databricks.sql.auth.common import ClientContext
 from databricks.sql.exc import SessionAlreadyClosedError, DatabaseError, RequestError
 from databricks.sql import __version__
@@ -96,10 +97,30 @@ def __init__(
         # Use the provided HTTP client (created in Connection)
         self.http_client = http_client
 
-        # Create auth provider with HTTP client context
-        self.auth_provider = get_python_sql_connector_auth_provider(
-            server_hostname, http_client=self.http_client, **kwargs
-        )
+        # Create auth provider with HTTP client context.
+        #
+        # On the kernel path the kernel owns the entire auth lifecycle
+        # (it acquires/refreshes OAuth tokens itself from the raw
+        # credentials — see the kernel auth bridge). We must NOT build
+        # the connector's own provider here: for OAuth it would eagerly
+        # run the U2M browser flow / M2M token exchange at connect()
+        # time (``get_auth_provider`` invokes ``_initial_get_token`` in
+        # the provider constructor), racing — and conflicting with — the
+        # kernel's auth before ``use_kernel`` is even consulted.
+        #
+        # So for ``use_kernel`` we hand the bridge only a minimal PAT
+        # provider when an ``access_token`` is present, and ``None``
+        # otherwise (OAuth M2M/U2M resolve purely from the raw kwargs
+        # the bridge reads). The Thrift / SEA backends are unchanged.
+        if kwargs.get("use_kernel", False):
+            access_token = kwargs.get("access_token")
+            self.auth_provider = (
+                AccessTokenAuthProvider(access_token) if access_token else None
+            )
+        else:
+            self.auth_provider = get_python_sql_connector_auth_provider(
+                server_hostname, http_client=self.http_client, **kwargs
+            )
 
         self.backend = self._create_backend(
             server_hostname,
@@ -140,11 +161,11 @@ def _create_backend(
             logger.debug("Creating kernel-backed client for use_kernel=True")
             # Forward the raw auth-relevant connect() kwargs so the
             # kernel auth bridge can build OAuth kwargs from the
-            # original credentials. The OAuth client secret is consumed
-            # and discarded during ``auth_provider`` construction, so it
-            # can only be read from these raw kwargs — not the built
-            # provider. These are kernel-only; the Thrift / SEA backends
-            # are unaffected.
+            # original credentials. On this path we intentionally did
+            # NOT build the connector's own OAuth provider (see __init__
+            # above), so these raw kwargs are the only source of the
+            # OAuth client id/secret. These are kernel-only; the Thrift
+            # / SEA backends are unaffected.
             kernel_auth_options = {
                 "auth_type": kwargs.get("auth_type"),
                 "oauth_client_id": kwargs.get("oauth_client_id"),

tests/e2e/test_kernel_tls.py

@@ -0,0 +1,140 @@
+"""E2E TLS tests for ``use_kernel=True``, exercised through an
+intercepting HTTPS proxy (mitmproxy) against a live Databricks
+workspace.
+
+This is the connector-level counterpart to the kernel repo's
+``v0_tls_e2e.rs``. The kernel test proves the raw TLS handshake against
+the kernel's ``TlsConfig`` directly; this one proves the *whole stack*:
+``sql.connect(use_kernel=True, _tls_*=...)`` → ``SSLOptions`` →
+``_kernel_tls_kwargs`` → pyo3 → kernel handshake.
+
+mitmproxy sits in front of the workspace and re-signs every TLS
+connection with its own CA. A client trusting only the system roots
+sees an untrusted cert and fails — unless we add mitmproxy's CA via
+``_tls_trusted_ca_file`` or disable validation via ``_tls_no_verify``.
+The three tests cover exactly those outcomes (reject / CA-trusted /
+no-verify).
+
+Traffic is routed through mitmproxy via the ``HTTPS_PROXY`` env var
+(reqwest honours it by default), since the connector's proxy surface
+isn't yet wired to the kernel. The CI job sets
+``HTTPS_PROXY=http://localhost:8080``.
+
+Skipped automatically when the kernel wheel, the live creds, or
+``MITMPROXY_CA_CERT`` are absent — so it's a no-op in the default test
+run and only fires in the ``kernel-e2e`` TLS CI job (or locally with
+mitmproxy up; see that workflow / the kernel repo's ``v0_tls_e2e.rs``
+header for the docker incantation).
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+import databricks.sql as sql
+from databricks.sql.exc import Error as DatabricksSqlError
+
+# Same real-wheel guard as test_kernel_backend.py: a fake
+# ``databricks_sql_kernel`` ModuleType injected by the unit tests has no
+# ``__file__``; only a compiled wheel does.
+_kernel_mod = pytest.importorskip(
+    "databricks_sql_kernel",
+    reason="use_kernel=True requires the databricks-sql-kernel package",
+)
+if not getattr(_kernel_mod, "__file__", None):
+    pytest.skip(
+        "databricks_sql_kernel is a test stub (no __file__); "
+        "install the real wheel to run kernel TLS e2e tests",
+        allow_module_level=True,
+    )
+
+_MITM_CA = os.getenv("MITMPROXY_CA_CERT")
+if not _MITM_CA:
+    pytest.skip(
+        "MITMPROXY_CA_CERT not set — TLS e2e runs only behind the "
+        "mitmproxy CI job (or a local mitmproxy)",
+        allow_module_level=True,
+    )
+
+
+def _env(key):
+    v = os.getenv(key)
+    return v if v else None
+
+
+@pytest.fixture(scope="module")
+def base_params(connection_details):
+    """Live workspace connection params for use_kernel=True. Prefers
+    OAuth M2M (so the TLS suite doubles as M2M-through-proxy coverage)
+    and falls back to PAT."""
+    host = connection_details.get("host")
+    http_path = connection_details.get("http_path")
+    if not (host and http_path):
+        pytest.skip("DATABRICKS_SERVER_HOSTNAME / DATABRICKS_HTTP_PATH not set")
+
+    params = {
+        "server_hostname": host,
+        "http_path": http_path,
+        "use_kernel": True,
+    }
+
+    client_id = _env("DATABRICKS_TEST_CLIENT_ID")
+    client_secret = _env("DATABRICKS_TEST_CLIENT_SECRET")
+    token = connection_details.get("access_token")
+    if client_id and client_secret:
+        params["oauth_client_id"] = client_id
+        params["oauth_client_secret"] = client_secret
+    elif token:
+        params["access_token"] = token
+    else:
+        pytest.skip(
+            "need DATABRICKS_TEST_CLIENT_ID/SECRET (OAuth M2M) or "
+            "DATABRICKS_TOKEN (PAT)"
+        )
+    return params
+
+
+def _select_one(conn):
+    with conn.cursor() as cur:
+        cur.execute("SELECT 1 AS n")
+        rows = cur.fetchall()
+    assert len(rows) == 1 and rows[0][0] == 1
+
+
+def test_tls_fails_without_config_behind_intercepting_proxy(base_params):
+    """Strict default (no CA, verify on) must reject mitmproxy's
+    re-signed certificate. Guards that the positive tests below are
+    actually validating the chain rather than bypassing the proxy."""
+    with pytest.raises(DatabricksSqlError):
+        conn = sql.connect(**base_params)
+        try:
+            # Some auth flows defer the first network call until a query.
+            _select_one(conn)
+        finally:
+            try:
+                conn.close()
+            except Exception:
+                pass
+
+
+def test_tls_with_trusted_custom_ca_succeeds(base_params):
+    """Adding mitmproxy's CA via ``_tls_trusted_ca_file`` makes the
+    re-signed cert trusted, so the full round-trip succeeds. Proves
+    ``_tls_trusted_ca_file`` is wired through to the kernel handshake."""
+    conn = sql.connect(_tls_trusted_ca_file=_MITM_CA, **base_params)
+    try:
+        _select_one(conn)
+    finally:
+        conn.close()
+
+
+def test_tls_with_no_verify_succeeds(base_params):
+    """``_tls_no_verify=True`` disables chain validation, so the
+    round-trip succeeds without supplying the CA."""
+    conn = sql.connect(_tls_no_verify=True, **base_params)
+    try:
+        _select_one(conn)
+    finally:
+        conn.close()