feat: add oauth_token_cache_enabled to control kernel U2M on-disk token cache

Adds the oauth_token_cache_enabled connect() kwarg, forwarded to the pyo3 Session token_cache_enabled field on the oauth-u2m path. U2M-only; disabled by default (matches Thrift no-persistence posture); enable-flag only; experimental_oauth_persistence is untouched.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

Author: eric-wang-1990

CONNECTION_PARAMETERS.md

@@ -82,6 +82,7 @@ to change without notice.
 | `credentials_provider`                              | `CredentialsProvider`|   ✅   |   ❌   | `None`                    | Custom external credentials provider. **Rejected on the kernel path** (`NotSupportedError`) — it is an opaque token source, so the kernel cannot own the token lifecycle; use `oauth_client_id` + `oauth_client_secret` for M2M, or the Thrift backend. |
 | `identity_federation_client_id`                     | `str`                |   ✅   |   ✅   | `None`                    | Workload identity / token-federation client id (kernel support added in #910).                                                                                 |
 | `experimental_oauth_persistence`                    | `OAuthPersistence`   |   ✅   |   ❌   | `None`                    | **Thrift-only.** The kernel owns its own token lifecycle and does not accept a persistence store.                                                              |
+| `oauth_token_cache_enabled`                         | `bool \| None`       |   ❌   |   ✅   | `None`                    | **Kernel-only, U2M-only.** Controls whether the kernel persists OAuth U2M refresh tokens to disk (AES-256 encrypted, at `~/.config/databricks-sql-kernel/oauth/`, requires databricks-sql-kernel PR #283). When unset (None, default), the kernel's default applies; False disables persistence (in-memory only); True enables on-disk cache. Distinct from `experimental_oauth_persistence` — this controls the kernel's built-in encrypted storage, not a pluggable callback.  |
 | `azure_client_id` / `azure_client_secret` / `azure_tenant_id` | `str` | ✅ | ✅ | `None` | Azure service-principal (Entra ID M2M), selected by `auth_type="azure-sp-m2m"`. On the kernel path the connector forwards these to the kernel, which owns Azure resolution (Entra v2.0 token endpoint + the Databricks-resource `.default` scope) (#919). **`azure_tenant_id` is optional on the kernel path too** — like Thrift, the kernel auto-discovers it from the workspace's `/aad/auth` redirect when omitted. |
 | `azure_workspace_resource_id`                       | `str`                |   ✅   |   ✅   | `None`                    | For `azure-sp-m2m`. When set, the SP **management token** (`X-Databricks-Azure-SP-Management-Token`) + `X-Databricks-Azure-Workspace-Resource-Id` header are sent, to authorize an SP that has an Azure RBAC role but is not a workspace member. Omit it for a workspace-member SP (the data token authenticates alone; no management token is fetched). Works on both the kernel and Thrift paths. |
 | `_use_cert_as_auth` (+ `_tls_client_cert_file`)     | `bool`               |   ✅   |   ❌   | `False`                   | Authenticate with a TLS client certificate instead of a token. Thrift-only.                                                                                    |

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

@@ -460,6 +460,12 @@ def kernel_auth_kwargs(
                 else list(PYSQL_OAUTH_REDIRECT_PORT_RANGE)
             ),
             "oauth_scopes": scopes if scopes is not None else list(PYSQL_OAUTH_SCOPES),
+            # OAuth U2M token-cache enable/disable: when present in auth_options,
+            # forward to the kernel as token_cache_enabled on the U2M branch.
+            # Default disabled (bool(None) = False) for backward compatibility when
+            # moving token persistence control to the kernel. This ensures callers
+            # must opt-in to on-disk persistence rather than silently enabling it.
+            "token_cache_enabled": bool(opts.get("oauth_token_cache_enabled")),
         }
         if federation_client_id:
             kwargs["identity_federation_client_id"] = federation_client_id

src/databricks/sql/client.py

@@ -225,6 +225,17 @@ def read(self) -> Optional[OAuthToken]:
                             experimental_oauth_persistence=DevOnlyFilePersistence("~/dev-oauth.json")
                         )
                 ```
+            :param oauth_token_cache_enabled: `bool | None`, optional (default is None)
+                **Kernel-only, U2M-only.** Controls whether the kernel persists OAuth U2M
+                refresh tokens to disk (AES-256 encrypted, at `~/.config/databricks-sql-kernel/oauth/`).
+                When unset (None, the default), the kernel's own default behavior applies.
+                When True, enables persistent on-disk token cache; when False, tokens are
+                held in memory only and the user must re-authenticate when the process restarts.
+                Has no effect on Thrift or SEA backends, which maintain their own token
+                lifecycle via `experimental_oauth_persistence`. This parameter is distinct
+                from the Thrift-only `experimental_oauth_persistence` — this controls the
+                kernel's built-in encrypted storage, whereas `experimental_oauth_persistence`
+                is a pluggable callback interface for Thrift-path custom storage.
             :param _use_arrow_native_complex_types: `bool`, optional
                 Controls whether a complex type field value is returned as a string or as a native Arrow type. Defaults to True.
                 When True:

src/databricks/sql/session.py

@@ -186,6 +186,12 @@ def _create_backend(
                 "identity_federation_client_id": kwargs.get(
                     "identity_federation_client_id"
                 ),
+                # OAuth U2M token-cache enable/disable: controls whether the kernel
+                # persists U2M refresh tokens to disk (encrypted, at ~/.config/databricks-sql-kernel/oauth/).
+                # Omitted ⇒ kernel default (enabled); False ⇒ in-memory only.
+                # This is forwarded to the kernel's pyo3 Session as token_cache_enabled
+                # on the oauth-u2m auth branch only, ensuring backward compat when moved to the kernel path.
+                "oauth_token_cache_enabled": kwargs.get("oauth_token_cache_enabled"),
                 # Azure Entra SP credentials for the azure-sp-m2m path. The
                 # kernel owns Azure resolution (endpoint/scope/tenant discovery),
                 # so these raw kwargs are the only source; without threading them

tests/unit/test_kernel_auth_bridge.py

@@ -432,6 +432,8 @@ def test_bare_databricks_oauth_forwards_full_python_bundle(self):
             # Full registered port list → the kernel binds the first free one.
             "redirect_ports": list(PYSQL_OAUTH_REDIRECT_PORT_RANGE),
             "oauth_scopes": list(PYSQL_OAUTH_SCOPES),
+            # token_cache_enabled defaults to False (disable-by-default).
+            "token_cache_enabled": False,
         }
 
     def test_azure_oauth_maps_to_in_house_u2m(self):
@@ -450,6 +452,8 @@ def test_azure_oauth_maps_to_in_house_u2m(self):
             "client_id": PYSQL_OAUTH_CLIENT_ID,
             "redirect_ports": list(PYSQL_OAUTH_REDIRECT_PORT_RANGE),
             "oauth_scopes": list(PYSQL_OAUTH_SCOPES),
+            # token_cache_enabled defaults to False (disable-by-default).
+            "token_cache_enabled": False,
         }
 
     def test_azure_oauth_honors_custom_client_id_port_and_scopes(self):
@@ -469,6 +473,8 @@ def test_azure_oauth_honors_custom_client_id_port_and_scopes(self):
             "client_id": "custom-client",
             "redirect_ports": [9999],
             "oauth_scopes": ["custom-scope", "offline_access"],
+            # token_cache_enabled defaults to False (disable-by-default).
+            "token_cache_enabled": False,
         }
 
     def test_u2m_custom_client_id_port_and_scopes_honored(self):
@@ -489,6 +495,8 @@ def test_u2m_custom_client_id_port_and_scopes_honored(self):
             "client_id": "custom-client",
             "redirect_ports": [9999],
             "oauth_scopes": ["custom-scope", "offline_access"],
+            # token_cache_enabled defaults to False (disable-by-default).
+            "token_cache_enabled": False,
         }
 
     def test_u2m_custom_client_id_only_falls_back_to_connector_defaults(self):
@@ -507,6 +515,8 @@ def test_u2m_custom_client_id_only_falls_back_to_connector_defaults(self):
             "client_id": "custom-client",
             "redirect_ports": list(PYSQL_OAUTH_REDIRECT_PORT_RANGE),
             "oauth_scopes": list(PYSQL_OAUTH_SCOPES),
+            # token_cache_enabled defaults to False (disable-by-default).
+            "token_cache_enabled": False,
         }
 
     def test_u2m_redirect_port_coerced_to_int(self):
@@ -578,6 +588,76 @@ def test_u2m_normalizes_space_delimited_scopes(self):
         )
         assert kwargs["oauth_scopes"] == ["all-apis", "offline_access"]
 
+    def test_u2m_token_cache_enabled_unset_defaults_to_false(self):
+        # When oauth_token_cache_enabled is omitted, the kernel U2M kwargs
+        # must include token_cache_enabled=False (disable-by-default) so the
+        # kernel does not silently start persisting tokens to disk.
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {"auth_type": "databricks-oauth"},
+        )
+        assert kwargs["token_cache_enabled"] is False
+
+    def test_u2m_token_cache_enabled_false_forwarded(self):
+        # When oauth_token_cache_enabled=False, forward token_cache_enabled=False.
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {
+                "auth_type": "databricks-oauth",
+                "oauth_token_cache_enabled": False,
+            },
+        )
+        assert kwargs["token_cache_enabled"] is False
+
+    def test_u2m_token_cache_enabled_true_forwarded(self):
+        # When oauth_token_cache_enabled=True, forward token_cache_enabled=True.
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {
+                "auth_type": "databricks-oauth",
+                "oauth_token_cache_enabled": True,
+            },
+        )
+        assert kwargs["token_cache_enabled"] is True
+
+    @pytest.mark.parametrize("u2m_auth_type", ["databricks-oauth", "azure-oauth"])
+    def test_u2m_token_cache_enabled_both_auth_types(self, u2m_auth_type):
+        # token_cache_enabled applies to both databricks-oauth and azure-oauth U2M types.
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {
+                "auth_type": u2m_auth_type,
+                "oauth_token_cache_enabled": True,
+            },
+        )
+        assert kwargs["token_cache_enabled"] is True
+
+    def test_token_cache_enabled_not_forwarded_to_m2m(self):
+        # oauth_token_cache_enabled should NOT be forwarded on the M2M path
+        # (M2M handles its own token lifecycle independently).
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {
+                "oauth_client_id": "sp-uuid",
+                "oauth_client_secret": "shh",
+                "oauth_token_cache_enabled": True,
+            },
+        )
+        # On the M2M path, token_cache_enabled should NOT be present.
+        assert "token_cache_enabled" not in kwargs
+        assert kwargs["auth_type"] == "oauth-m2m"
+
+    def test_token_cache_enabled_not_forwarded_to_pat(self):
+        # oauth_token_cache_enabled should NOT be forwarded on the PAT path
+        # (PAT is a static token with no refresh/cache mechanism).
+        kwargs = kernel_auth_kwargs(
+            AccessTokenAuthProvider("dapi-xyz"),
+            {"oauth_token_cache_enabled": True},
+        )
+        # On the PAT path, token_cache_enabled should NOT be present.
+        assert "token_cache_enabled" not in kwargs
+        assert kwargs["auth_type"] == "pat"
+
 
 class TestKernelIdentityFederationClientId:
     @pytest.mark.parametrize(