Skip to content

Commit aab4203

Browse files
ai: apply changes for #932 (1 review thread)
Addresses: - #3859712688 at src/databricks/sql/session.py:197 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent aadfdd2 commit aab4203

1 file changed

Lines changed: 42 additions & 0 deletions

File tree

tests/unit/test_session.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,48 @@ def test_azure_sp_m2m_kwargs_threaded_into_kernel_auth_options(self):
526526
finally:
527527
conn.close()
528528

529+
def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(self):
530+
# oauth_token_cache_enabled must reach the kernel auth bridge via
531+
# auth_options; without this session.py mapping line the feature
532+
# would silently regress to always-disabled (the safe default masks
533+
# the failure). Guards the session.py -> kernel_auth_options map.
534+
import sys
535+
import types
536+
537+
pytest.importorskip(
538+
"pyarrow",
539+
reason="kernel client module imports pyarrow at load",
540+
)
541+
542+
fake = types.ModuleType("databricks_sql_kernel")
543+
fake.KernelError = type("KernelError", (Exception,), {})
544+
fake.Session = MagicMock()
545+
546+
with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
547+
"databricks.sql.backend.kernel.client.KernelDatabricksClient"
548+
) as mock_kernel_client, patch(
549+
"%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
550+
):
551+
instance = mock_kernel_client.return_value
552+
instance.open_session.return_value = SessionId(
553+
BackendType.SEA, "sess-id", None
554+
)
555+
556+
conn = databricks.sql.connect(
557+
server_hostname="foo",
558+
http_path="/sql/1.0/warehouses/abc",
559+
use_kernel=True,
560+
auth_type="databricks-oauth",
561+
oauth_token_cache_enabled=True,
562+
enable_telemetry=False,
563+
)
564+
try:
565+
_, kwargs = mock_kernel_client.call_args
566+
opts = kwargs["auth_options"]
567+
assert opts["oauth_token_cache_enabled"] is True
568+
finally:
569+
conn.close()
570+
529571

530572
class TestKernelUserAgentForwarding:
531573
"""user_agent_entry must reach the kernel on the use_kernel path —

0 commit comments

Comments
 (0)