From e2e33b413c90109bfa15dddde1e220be804ac0f2 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 25 Aug 2026 16:18:07 -0700 Subject: [PATCH 1/8] 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 --- CONNECTION_PARAMETERS.md | 1 + .../sql/backend/kernel/auth_bridge.py | 6 ++ src/databricks/sql/client.py | 11 +++ src/databricks/sql/session.py | 6 ++ tests/unit/test_kernel_auth_bridge.py | 80 +++++++++++++++++++ 5 files changed, 104 insertions(+) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index f63de23a9..9a18033e3 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, 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. | diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 1fed54c41..8ee84d672 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/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 diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 44895954f..ed0eb2250 100755 --- a/src/databricks/sql/client.py +++ b/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: diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index f35cdf525..99ebba24c 100644 --- a/src/databricks/sql/session.py +++ b/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 diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index abce5a0f4..ddf6c1c91 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,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( From 0de2e82ab6c0856884e78c9dcadba5f9498a7448 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 25 Aug 2026 16:50:36 -0700 Subject: [PATCH 2/8] docs: correct oauth_token_cache_enabled default to disabled in CONNECTION_PARAMETERS.md The connector forwards bool(None)=False on unset, so omitting the kwarg disables the on-disk cache; it does not inherit the kernel's enabled-by-default. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 --- CONNECTION_PARAMETERS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 9a18033e3..98804f495 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -82,7 +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. | +| `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). **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. | From baf88f00e37d9513b3848754296eb4ee6fc4a2f5 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Wed, 26 Aug 2026 00:19:09 +0000 Subject: [PATCH 3/8] ai: apply changes for #932 (3 review threads) Addresses: - #3858601138 at src/databricks/sql/backend/kernel/auth_bridge.py:468 - #3858601153 at src/databricks/sql/session.py:193 - #3858601164 at src/databricks/sql/client.py:234 Signed-off-by: peco-engineer-bot[bot] --- .../sql/backend/kernel/auth_bridge.py | 34 ++++++++++++++++--- src/databricks/sql/client.py | 10 ++++-- src/databricks/sql/session.py | 7 ++-- tests/unit/test_kernel_auth_bridge.py | 29 ++++++++++++++++ 4 files changed, 71 insertions(+), 9 deletions(-) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 8ee84d672..4cb874654 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -462,10 +462,14 @@ def kernel_auth_kwargs( "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")), + # Default disabled 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. + # Coerced via _coerce_bool so a string DSN/env value like "False" + # is not treated as truthy (bool("False") is True). + "token_cache_enabled": _coerce_bool( + opts.get("oauth_token_cache_enabled") + ), } if federation_client_id: kwargs["identity_federation_client_id"] = federation_client_id @@ -519,6 +523,28 @@ def _coerce_redirect_port(redirect_port: Any) -> int: ) +def _coerce_bool(value: Any) -> bool: + """Coerce an opt-in boolean flag (e.g. ``oauth_token_cache_enabled``, + which may arrive as a string from a DSN/env) to a ``bool``. + + A plain ``bool(value)`` is wrong for string inputs: ``bool("False")`` is + ``True``, which would silently enable on-disk token persistence whenever + the flag arrived as the string ``"False"``. Only genuinely truthy values + enable the flag: real booleans, and the usual textual/numeric truthy + spellings ("true"/"1"/"yes"/"on"). ``None`` (unset) and anything else + disable it (opt-in default).""" + if isinstance(value, bool): + return value + if value is None: + return False + if isinstance(value, str): + return value.strip().lower() in ("true", "1", "yes", "on") + if isinstance(value, (int, float)): + return value != 0 + # Unknown types default to disabled rather than truthy-by-accident. + return False + + def _normalize_scopes(scopes: Any) -> Optional[list]: """Normalise an ``oauth_scopes`` value to a list of strings, or ``None`` to let the kernel apply its defaults. diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index ed0eb2250..8ba5e3f8c 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -228,9 +228,13 @@ def read(self) -> Optional[OAuthToken]: :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. + 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 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 diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 99ebba24c..4be824202 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -188,9 +188,12 @@ def _create_backend( ), # 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. + # Coerced via _coerce_bool on the oauth-u2m branch, so 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, ensuring backward compat when moved to the kernel path. + # 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), diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index ddf6c1c91..d0386004d 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -620,6 +620,35 @@ def test_u2m_token_cache_enabled_true_forwarded(self): ) assert kwargs["token_cache_enabled"] is True + @pytest.mark.parametrize( + "raw_value", + ["False", "false", "0", "no", "off", "", " ", "nope"], + ) + def test_u2m_token_cache_enabled_falsey_string_stays_false(self, raw_value): + # A string DSN/env value that reads as falsey (e.g. "False") must NOT + # enable on-disk persistence: bool("False") is True, so the flag is + # coerced via _coerce_bool rather than bool(). + 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("raw_value", ["True", "true", "1", "yes", "on"]) + def test_u2m_token_cache_enabled_truthy_string_enables(self, raw_value): + # An explicit truthy string DSN/env value enables persistence. + kwargs = kernel_auth_kwargs( + _FakeOAuthProvider(), + { + "auth_type": "databricks-oauth", + "oauth_token_cache_enabled": raw_value, + }, + ) + 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. From aadfdd2ff3e8c59830f69ebded0e588e6428b789 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Wed, 26 Aug 2026 04:39:14 +0000 Subject: [PATCH 4/8] ai: apply changes for #932 (1 review thread) Addresses: - #3859272243 at src/databricks/sql/client.py:238 Signed-off-by: peco-engineer-bot[bot] --- src/databricks/sql/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 8ba5e3f8c..550e79ec6 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -235,7 +235,7 @@ def read(self) -> Optional[OAuthToken]: 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 Thrift or SEA backends, which maintain their own token + 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` From aab4203ba068e0dd53bda7fad88cb3f19950ca80 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Wed, 26 Aug 2026 04:47:26 +0000 Subject: [PATCH 5/8] ai: apply changes for #932 (1 review thread) Addresses: - #3859712688 at src/databricks/sql/session.py:197 Signed-off-by: peco-engineer-bot[bot] --- tests/unit/test_session.py | 42 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 6fcefcade..aad3acf96 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -526,6 +526,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 — From 4179dd58d02bde40f39dda3b9290cb8af7d2ea25 Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 25 Aug 2026 22:41:49 -0700 Subject: [PATCH 6/8] refactor(kernel-auth): drop _coerce_bool; treat oauth_token_cache_enabled as Optional[bool] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit oauth_token_cache_enabled is a typed connection kwarg like the connector's other booleans (use_cloud_fetch, _use_arrow_native_complex_types, ...), none of which coerce string inputs. Replace the one-off _coerce_bool with `opts.get("oauth_token_cache_enabled") is True`: only a real True enables the kernel's on-disk U2M token cache, and unset (None)/False forward an explicit False so the kernel — whose own default is enabled — stays disabled by default. Any non-bool value fails safe to disabled rather than silently enabling. Matches the nodejs connector's `tokenCacheEnabled ?? false` and removes the inconsistency of coercing this one flag while every sibling boolean is passed raw. Tests: the string-coercion cases are replaced by one asserting non-bool inputs never enable the cache. Signed-off-by: eric-wang-1990 --- .../sql/backend/kernel/auth_bridge.py | 40 ++++--------------- src/databricks/sql/session.py | 2 +- tests/unit/test_kernel_auth_bridge.py | 23 +++-------- 3 files changed, 15 insertions(+), 50 deletions(-) diff --git a/src/databricks/sql/backend/kernel/auth_bridge.py b/src/databricks/sql/backend/kernel/auth_bridge.py index 4cb874654..3477caf6b 100644 --- a/src/databricks/sql/backend/kernel/auth_bridge.py +++ b/src/databricks/sql/backend/kernel/auth_bridge.py @@ -460,16 +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 token-cache enable/disable: when present in auth_options, - # forward to the kernel as token_cache_enabled on the U2M branch. - # Default disabled 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. - # Coerced via _coerce_bool so a string DSN/env value like "False" - # is not treated as truthy (bool("False") is True). - "token_cache_enabled": _coerce_bool( - opts.get("oauth_token_cache_enabled") - ), + # 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 @@ -523,28 +521,6 @@ def _coerce_redirect_port(redirect_port: Any) -> int: ) -def _coerce_bool(value: Any) -> bool: - """Coerce an opt-in boolean flag (e.g. ``oauth_token_cache_enabled``, - which may arrive as a string from a DSN/env) to a ``bool``. - - A plain ``bool(value)`` is wrong for string inputs: ``bool("False")`` is - ``True``, which would silently enable on-disk token persistence whenever - the flag arrived as the string ``"False"``. Only genuinely truthy values - enable the flag: real booleans, and the usual textual/numeric truthy - spellings ("true"/"1"/"yes"/"on"). ``None`` (unset) and anything else - disable it (opt-in default).""" - if isinstance(value, bool): - return value - if value is None: - return False - if isinstance(value, str): - return value.strip().lower() in ("true", "1", "yes", "on") - if isinstance(value, (int, float)): - return value != 0 - # Unknown types default to disabled rather than truthy-by-accident. - return False - - def _normalize_scopes(scopes: Any) -> Optional[list]: """Normalise an ``oauth_scopes`` value to a list of strings, or ``None`` to let the kernel apply its defaults. diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 4be824202..f806663b2 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -188,7 +188,7 @@ def _create_backend( ), # OAuth U2M token-cache enable/disable: controls whether the kernel # persists U2M refresh tokens to disk (encrypted, at ~/.config/databricks-sql-kernel/oauth/). - # Coerced via _coerce_bool on the oauth-u2m branch, so omitted/None + # 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. diff --git a/tests/unit/test_kernel_auth_bridge.py b/tests/unit/test_kernel_auth_bridge.py index d0386004d..000515915 100644 --- a/tests/unit/test_kernel_auth_bridge.py +++ b/tests/unit/test_kernel_auth_bridge.py @@ -622,12 +622,13 @@ def test_u2m_token_cache_enabled_true_forwarded(self): @pytest.mark.parametrize( "raw_value", - ["False", "false", "0", "no", "off", "", " ", "nope"], + ["True", "true", "1", "yes", "on", "False", "false", "0", "", 1, 0], ) - def test_u2m_token_cache_enabled_falsey_string_stays_false(self, raw_value): - # A string DSN/env value that reads as falsey (e.g. "False") must NOT - # enable on-disk persistence: bool("False") is True, so the flag is - # coerced via _coerce_bool rather than bool(). + 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(), { @@ -637,18 +638,6 @@ def test_u2m_token_cache_enabled_falsey_string_stays_false(self, raw_value): ) assert kwargs["token_cache_enabled"] is False - @pytest.mark.parametrize("raw_value", ["True", "true", "1", "yes", "on"]) - def test_u2m_token_cache_enabled_truthy_string_enables(self, raw_value): - # An explicit truthy string DSN/env value enables persistence. - kwargs = kernel_auth_kwargs( - _FakeOAuthProvider(), - { - "auth_type": "databricks-oauth", - "oauth_token_cache_enabled": raw_value, - }, - ) - 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. From 6de668aecc843ae1112ee8fd66eb3193f1c7597e Mon Sep 17 00:00:00 2001 From: eric-wang-1990 Date: Tue, 25 Aug 2026 23:45:56 -0700 Subject: [PATCH 7/8] chore(kernel): bump KERNEL_REV to 628abd6 (kernel #283, token_cache_enabled) The token_cache_enabled field this PR forwards to the kernel exists in the pyo3 Session only as of kernel #283. The prior pin (ad78a5b) lacks it, so the flag reached a kernel that ignored it. Bump to 628abd6 (kernel main HEAD): its pyo3 Session signature is a backward-compatible superset of ad78a5b's (identical azure-SP params, plus token_cache_enabled/_passphrase), so the kernel-e2e CI builds a kernel that honors the flag. Verified E2E (U2M): unset => no on-disk cache; True => encrypted cache file written. Published databricks-sql-kernel pin stays ^0.2.0 (no 0.2.1 pyo3 wheel yet); KERNEL_REV is what the kernel-e2e workflow builds the wheel from. Signed-off-by: eric-wang-1990 --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index 6cd3da53d..8d29e5090 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -ad78a5be3dc8bb7fc78ec574492515ab24e23d4c +628abd6f5045897efcadb38ec77a1e9e0c23544e From 7a8d315d0cb15239beb3cb698a70adb1f0ba1261 Mon Sep 17 00:00:00 2001 From: "peco-engineer-bot[bot]" Date: Wed, 26 Aug 2026 06:57:02 +0000 Subject: [PATCH 8/8] ai: apply changes for #932 (1 review thread) Addresses: - #3860320318 at src/databricks/sql/client.py:230 Signed-off-by: peco-engineer-bot[bot] --- CONNECTION_PARAMETERS.md | 2 +- src/databricks/sql/client.py | 4 +++- src/databricks/sql/session.py | 4 +++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index b81e3e48e..599ab4fb4 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -82,7 +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). **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. | +| `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/src/databricks/sql/client.py b/src/databricks/sql/client.py index 17dae9d2a..e34999b18 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -227,7 +227,9 @@ def read(self) -> Optional[OAuthToken]: ``` :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/`). + 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 diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 27410d893..1c7f74bf6 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -187,7 +187,9 @@ def _create_backend( "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/). + # 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