Mirror the Thrift auth default onto the Reyden kernel recovery path

On auto-recovery the Thrift default of auth_type=None (implicitly
databricks-oauth) was forwarded unchanged to the kernel, which has no
such fallback and rejects auth_type=None unless a PAT/M2M credential
shape is present — so a bare OAuth-U2M connection failed to recover.
Inject auth_type=databricks-oauth on the default-path kernel retry when
auth_type is unset and no credential shape is present, mirroring Thrift;
skip it when a credential shape exists so kernel routing is unchanged.

Also correct the warehouse-cache docstring: warehouse ids are globally
unique, so there is no cross-workspace collision on a shared SPOG host;
the host component is an optimization, not a correctness requirement.

Co-authored-by: Isaac <no-reply@databricks.com>
Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>

Author: rahuls-db

src/databricks/sql/backend/reyden_warehouse_cache.py

@@ -6,10 +6,13 @@
 here so later connections to the same warehouse skip the doomed Thrift attempt
 and open on the kernel directly.
 
-Keyed by ``(host, warehouse_id)`` — the host is part of the key so the same
-warehouse id observed on two different workspaces never collides. Entries
-expire after ``_TTL_SECONDS`` so a warehouse later reconfigured to accept Thrift
-is eventually retried.
+Keyed by ``(host, warehouse_id)``. Warehouse ids are globally unique, so the
+warehouse id alone identifies the warehouse — even on a SPOG host shared by many
+workspaces (where only the ``?o=<workspace-id>`` path param distinguishes them),
+there is no cross-workspace collision. The host is kept in the key only as a
+cheap optimization (scoping lookups) and defense-in-depth, not for correctness.
+Entries expire after ``_TTL_SECONDS`` so a warehouse later reconfigured to accept
+Thrift is eventually retried.
 """
 
 import re

src/databricks/sql/client.py

@@ -72,7 +72,7 @@
 from databricks.sql.session import Session
 from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId
 
-from databricks.sql.auth.common import ClientContext
+from databricks.sql.auth.common import AuthType, ClientContext
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.http import HttpMethod
 
@@ -563,6 +563,27 @@ def build_session(session_kwargs: dict) -> Session:
             self.session.open()
             return self.session
 
+        def kernel_recovery_kwargs() -> dict:
+            # Kwargs for re-opening on the kernel. The Thrift path treats an
+            # unset auth_type as databricks-oauth (see get_auth_provider); the
+            # kernel path has no such fallback and rejects auth_type=None unless
+            # a credential shape (PAT / OAuth M2M) is present. Mirror the Thrift
+            # default so a bare OAuth-U2M connection recovers instead of failing
+            # with NotSupportedError. Skip the injection when a credential shape
+            # is already present — the kernel routes on it regardless of
+            # auth_type, and forcing databricks-oauth alongside an M2M secret or
+            # a credentials_provider would change that routing.
+            recovery_kwargs = {**kwargs, "use_kernel": True}
+            has_credential_shape = (
+                recovery_kwargs.get("access_token")
+                or recovery_kwargs.get("oauth_client_secret")
+                or recovery_kwargs.get("oauth_jwt_key_file")
+                or recovery_kwargs.get("credentials_provider")
+            )
+            if recovery_kwargs.get("auth_type") is None and not has_credential_shape:
+                recovery_kwargs["auth_type"] = AuthType.DATABRICKS_OAUTH.value
+            return recovery_kwargs
+
         # An explicit backend choice is always honored — auto-recovery engages
         # only on the default (Thrift) path.
         explicit_backend = kwargs.get("use_kernel", False) or kwargs.get(
@@ -582,7 +603,7 @@ def build_session(session_kwargs: dict) -> Session:
                 warehouse_id,
                 server_hostname,
             )
-            return build_session({**kwargs, "use_kernel": True})
+            return build_session(kernel_recovery_kwargs())
 
         try:
             return build_session(kwargs)
@@ -597,7 +618,7 @@ def build_session(session_kwargs: dict) -> Session:
             if warehouse_id:
                 mark_reyden(server_hostname, warehouse_id)
             try:
-                return build_session({**kwargs, "use_kernel": True})
+                return build_session(kernel_recovery_kwargs())
             except Exception as kernel_ex:
                 # Surface the kernel failure (the actionable one) while keeping
                 # the original Thrift rejection in the chain for diagnosis.

tests/unit/test_session.py

@@ -930,3 +930,38 @@ def test_recovered_kernel_failure_suppresses_wrapper_telemetry(self, mock_thrift
                 self._connect(enable_telemetry=True)
         mock_fail_log.assert_called_once()
         assert mock_fail_log.call_args.kwargs["enable_telemetry"] is False
+
+    @patch("%s.session.get_python_sql_connector_auth_provider" % PACKAGE)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE)
+    def test_oauth_default_recovery_injects_databricks_oauth_auth_type(
+        self, mock_thrift, mock_provider
+    ):
+        # A bare connection (no access_token, no auth_type) defaults to OAuth
+        # U2M on the Thrift path. The kernel path has no such default and would
+        # reject auth_type=None, so recovery must inject databricks-oauth to
+        # mirror the Thrift default. (The provider builder is patched so the
+        # token-less Thrift attempt doesn't build a real OAuth provider, which
+        # would hit the network at construction.)
+        mock_thrift.return_value.open_session.side_effect = self._reject()
+        with self._fake_kernel() as mock_kernel:
+            conn = self._connect(access_token=None)
+            try:
+                _, kwargs = mock_kernel.call_args
+                assert kwargs["auth_options"]["auth_type"] == "databricks-oauth"
+            finally:
+                conn.close()
+
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE)
+    def test_pat_recovery_does_not_inject_auth_type(self, mock_thrift):
+        # With a credential shape present (here a PAT) the kernel routes on it
+        # regardless of auth_type, so recovery must NOT inject databricks-oauth:
+        # forcing it alongside other credentials could change kernel routing or
+        # trip the auth bridge's ambiguity guards.
+        mock_thrift.return_value.open_session.side_effect = self._reject()
+        with self._fake_kernel() as mock_kernel:
+            conn = self._connect()  # access_token="tok"
+            try:
+                _, kwargs = mock_kernel.call_args
+                assert kwargs["auth_options"].get("auth_type") is None
+            finally:
+                conn.close()