diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index b2f429a20..599ab4fb4 100644 --- a/CONNECTION_PARAMETERS.md +++ b/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. | diff --git a/KERNEL_REV b/KERNEL_REV index f751c496b..8d29e5090 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -dd810d6d0a179886b923c6e22dc785ddca16ebef +628abd6f5045897efcadb38ec77a1e9e0c23544e diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 1fed54c41..3477caf6b 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/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 diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 914a24dde..e34999b18 100755 --- a/src/databricks/sql/client.py +++ b/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: diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 19cd1dba3..1c7f74bf6 100644 --- a/src/databricks/sql/session.py +++ b/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 diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index abce5a0f4..000515915 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/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( diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index c50650e4d..2da8bd194 100644 --- a/tests/unit/test_session.py +++ b/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 —