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

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, in the OS config dir — `~/Library/Application Support/databricks-sql-kernel/oauth/` on macOS, `~/.config/databricks-sql-kernel/oauth/` on Linux; requires databricks-sql-kernel PR #283). **Disabled by default:** when unset (None) or False, the connector disables on-disk persistence (tokens in-memory only, matching Thrift); True enables the cache. Omitting it does **not** inherit the kernel's enabled-by-default. Distinct from `experimental_oauth_persistence` — this toggles 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.                                                                                    |

KERNEL_REV

@@ -1 +1 @@
-dd810d6d0a179886b923c6e22dc785ddca16ebef
+628abd6f5045897efcadb38ec77a1e9e0c23544e

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

@@ -460,6 +460,14 @@ 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 on-disk token cache. A typed Optional[bool], like the
+            # connector's other boolean options; only a real ``True`` enables
+            # it. The kernel's own default is *enabled*, so unset (None) must be
+            # forwarded as an explicit ``False`` — disabled, in-memory only —
+            # matching the Thrift posture so moving persistence control to the
+            # kernel never silently starts writing tokens to disk. Opt-in is
+            # therefore an explicit ``oauth_token_cache_enabled=True``.
+            "token_cache_enabled": opts.get("oauth_token_cache_enabled") is True,
         }
         if federation_client_id:
             kwargs["identity_federation_client_id"] = federation_client_id

src/databricks/sql/client.py

@@ -225,6 +225,23 @@ 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). The cache lives in the
+                OS config directory: `~/Library/Application Support/databricks-sql-kernel/oauth/`
+                on macOS, `~/.config/databricks-sql-kernel/oauth/` on Linux.
+                When unset (None, the default), the connector treats this as False and
+                forwards `token_cache_enabled=False` to the kernel, so on-disk caching is
+                disabled by default — matching the Thrift posture and avoiding silently
+                writing tokens to disk. Callers must opt in explicitly to enable persistence.
+                When True, enables persistent on-disk token cache; when False (or unset),
+                tokens are held in memory only and the user must re-authenticate when the
+                process restarts.
+                Has no effect on the Thrift backend, which maintains its 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,17 @@ 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, in the OS config dir:
+                # ~/Library/Application Support/databricks-sql-kernel/oauth/ on macOS,
+                # ~/.config/databricks-sql-kernel/oauth/ on Linux).
+                # A typed Optional[bool]; on the oauth-u2m branch omitted/None
+                # ⇒ token_cache_enabled=False (disabled, in-memory only) — the
+                # opt-in default that preserves backward compat when token
+                # persistence moves to the kernel path; True ⇒ on-disk persistence.
+                # This is forwarded to the kernel's pyo3 Session as token_cache_enabled
+                # on the oauth-u2m auth branch only.
+                "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,94 @@ 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(
+        "raw_value",
+        ["True", "true", "1", "yes", "on", "False", "false", "0", "", 1, 0],
+    )
+    def test_u2m_token_cache_enabled_non_bool_never_enables(self, raw_value):
+        # oauth_token_cache_enabled is a typed Optional[bool] (like the
+        # connector's other boolean options); only a real ``True`` enables
+        # on-disk persistence. Any non-bool value (e.g. a stray string from a
+        # DSN) fails safe to disabled rather than silently enabling.
+        kwargs = kernel_auth_kwargs(
+            _FakeOAuthProvider(),
+            {
+                "auth_type": "databricks-oauth",
+                "oauth_token_cache_enabled": raw_value,
+            },
+        )
+        assert kwargs["token_cache_enabled"] is False
+
+    @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(

tests/unit/test_session.py

@@ -527,6 +527,48 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self):
             finally:
                 conn.close()
 
+    def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(self):
+        # oauth_token_cache_enabled must reach the kernel auth bridge via
+        # auth_options; without this session.py mapping line the feature
+        # would silently regress to always-disabled (the safe default masks
+        # the failure). Guards the session.py -> kernel_auth_options map.
+        import sys
+        import types
+
+        pytest.importorskip(
+            "pyarrow",
+            reason="kernel client module imports pyarrow at load",
+        )
+
+        fake = types.ModuleType("databricks_sql_kernel")
+        fake.KernelError = type("KernelError", (Exception,), {})
+        fake.Session = MagicMock()
+
+        with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
+            "databricks.sql.backend.kernel.client.KernelDatabricksClient"
+        ) as mock_kernel_client, patch(
+            "%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
+        ):
+            instance = mock_kernel_client.return_value
+            instance.open_session.return_value = SessionId(
+                BackendType.SEA, "sess-id", None
+            )
+
+            conn = databricks.sql.connect(
+                server_hostname="foo",
+                http_path="/sql/1.0/warehouses/abc",
+                use_kernel=True,
+                auth_type="databricks-oauth",
+                oauth_token_cache_enabled=True,
+                enable_telemetry=False,
+            )
+            try:
+                _, kwargs = mock_kernel_client.call_args
+                opts = kwargs["auth_options"]
+                assert opts["oauth_token_cache_enabled"] is True
+            finally:
+                conn.close()
+
 
 class TestKernelUserAgentForwarding:
     """user_agent_entry must reach the kernel on the use_kernel path —