From aadb5e00aca4569db1ee8eea5cd7b5c4198f795a Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:08:51 +0000 Subject: [PATCH 1/7] perf: share SSL contexts across connections Each Thrift connect() loaded the CA bundle three to four times: in the connection's UnifiedHttpClient, in an SSL context THttpClient built but never used, in urllib3's per-TCP-connection context for the Thrift pool, and in the telemetry client's UnifiedHttpClient. Each load takes ~13 ms on OpenSSL 3.0 and does not scale across threads, which capped a process at ~15 connects/s regardless of thread count. Contexts are now cached process-wide, keyed by the TLS options and the CA files' stat signatures. The Thrift pool gets the same context urllib3 would have built, so TLS behavior is unchanged. Configurations with a client certificate are not cached. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- CHANGELOG.md | 1 + src/databricks/sql/auth/thrift_http_client.py | 46 +++- .../sql/common/ssl_context_cache.py | 91 +++++++ .../sql/common/unified_http_client.py | 47 ++-- tests/unit/conftest.py | 10 + tests/unit/test_ssl_context_cache.py | 243 ++++++++++++++++++ tests/unit/test_thrift_backend.py | 5 +- 7 files changed, 408 insertions(+), 35 deletions(-) create mode 100644 src/databricks/sql/common/ssl_context_cache.py create mode 100644 tests/unit/test_ssl_context_cache.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d9c58ebf3..1a148858d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ # Unreleased - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. +- Speed up `sql.connect()`, especially when many threads connect at once: connections with the same TLS settings now share one SSL context instead of reloading the CA bundle several times per connection. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 72a20bbf6..5236b9770 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -14,8 +14,10 @@ from thrift.transport.TTransport import CReadableTransport from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager from urllib3.util import make_headers +from urllib3.util.ssl_ import create_urllib3_context from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy from databricks.sql.types import SSLOptions +from databricks.sql.common import ssl_context_cache from databricks.sql.common.http_utils import ( detect_and_parse_proxy, ) @@ -23,6 +25,16 @@ logger = logging.getLogger(__name__) +def _create_urllib3_ssl_context(cert_reqs: int, ca_file: Optional[str]): + # The same context urllib3 builds per connection from cert_reqs/ca_certs. + context = create_urllib3_context(cert_reqs=cert_reqs) + if ca_file: + context.load_verify_locations(ca_file) + elif hasattr(context, "load_default_certs"): + context.load_default_certs() + return context + + class THttpClient(thrift.transport.THttpClient.THttpClient, CReadableTransport): realhost: Optional[str] realport: Optional[int] @@ -59,10 +71,8 @@ def __init__( assert self.scheme in ("http", "https") if self.scheme == "https": if self._ssl_options is not None: - # TODO: Not sure if those options are used anywhere - need to double-check self.certfile = self._ssl_options.tls_client_cert_file self.keyfile = self._ssl_options.tls_client_cert_key_file - self.context = self._ssl_options.create_ssl_context() self.port = parsed.port self.host = parsed.hostname self.path = parsed.path @@ -122,19 +132,27 @@ def open(self): pool_class = HTTPConnectionPool elif self.scheme == "https": pool_class = HTTPSConnectionPool - _pool_kwargs.update( - { - "cert_reqs": ( - ssl.CERT_REQUIRED - if self._ssl_options.tls_verify - else ssl.CERT_NONE - ), - "ca_certs": self._ssl_options.tls_trusted_ca_file, - "cert_file": self._ssl_options.tls_client_cert_file, - "key_file": self._ssl_options.tls_client_cert_key_file, - "key_password": self._ssl_options.tls_client_cert_key_password, - } + cert_reqs = ( + ssl.CERT_REQUIRED if self._ssl_options.tls_verify else ssl.CERT_NONE ) + _pool_kwargs["cert_reqs"] = cert_reqs + if ssl_context_cache.is_cacheable(self._ssl_options): + # Otherwise urllib3 builds and loads a new context per TCP connection. + ca_file = self._ssl_options.tls_trusted_ca_file + _pool_kwargs["ssl_context"] = ssl_context_cache.get_ssl_context( + "urllib3", + self._ssl_options, + lambda: _create_urllib3_ssl_context(cert_reqs, ca_file), + ) + else: + _pool_kwargs.update( + { + "ca_certs": self._ssl_options.tls_trusted_ca_file, + "cert_file": self._ssl_options.tls_client_cert_file, + "key_file": self._ssl_options.tls_client_cert_key_file, + "key_password": self._ssl_options.tls_client_cert_key_password, + } + ) if self.using_proxy(): proxy_manager = ProxyManager( diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py new file mode 100644 index 000000000..160f88e5e --- /dev/null +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -0,0 +1,91 @@ +"""Process-wide reuse of client SSL contexts. + +Loading a CA bundle into a new context costs ~10-15 ms with OpenSSL 3 and does +not scale across threads, so building one per connection caps how fast a +process can open connections. A cached context is reused only by callers with +identical TLS options and unchanged CA files. Configurations with a client +certificate are never cached. +""" + +import os +import ssl +import threading +from collections import OrderedDict +from typing import TYPE_CHECKING, Callable, Optional, Tuple + +if TYPE_CHECKING: + from databricks.sql.types import SSLOptions + +_MAX_ENTRIES = 32 + +_lock = threading.Lock() +_contexts: "OrderedDict[tuple, ssl.SSLContext]" = OrderedDict() + + +def _file_signature(path) -> Optional[Tuple]: + # Rewriting or replacing a CA file yields a new key, so rotation is picked up. + if not path: + return None + try: + st = os.stat(path) + except (OSError, TypeError, ValueError): + return (path, None) + return (path, st.st_ino, st.st_size, st.st_mtime_ns) + + +def _cache_key(kind: str, ssl_options: "SSLOptions") -> tuple: + defaults = ssl.get_default_verify_paths() # honors SSL_CERT_FILE / SSL_CERT_DIR + return ( + kind, + bool(ssl_options.tls_verify), + bool(ssl_options.tls_verify_hostname), + _file_signature(ssl_options.tls_trusted_ca_file), + _file_signature(defaults.cafile), + defaults.capath, + os.environ.get("SSLKEYLOGFILE"), + ) + + +def is_cacheable(ssl_options: "SSLOptions") -> bool: + return not ( + ssl_options.tls_client_cert_file or ssl_options.tls_client_cert_key_file + ) + + +def get_ssl_context( + kind: str, ssl_options: "SSLOptions", build: Callable[[], ssl.SSLContext] +) -> ssl.SSLContext: + """Return the shared context for ``kind`` and ``ssl_options``, building it once. + + ``kind`` separates builders that configure contexts differently. Callers + must not modify the returned context. + """ + if not is_cacheable(ssl_options): + return build() + + key = _cache_key(kind, ssl_options) + with _lock: + context = _contexts.get(key) + if context is None: + context = build() + _contexts[key] = context + if len(_contexts) > _MAX_ENTRIES: + _contexts.popitem(last=False) + else: + _contexts.move_to_end(key) + return context + + +def clear() -> None: + with _lock: + _contexts.clear() + + +def _reset_after_fork() -> None: + # A fork during a build would leave the child's copy of the lock held. + global _lock + _lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_after_fork) diff --git a/src/databricks/sql/common/unified_http_client.py b/src/databricks/sql/common/unified_http_client.py index becf2eef3..90a92307f 100644 --- a/src/databricks/sql/common/unified_http_client.py +++ b/src/databricks/sql/common/unified_http_client.py @@ -20,6 +20,7 @@ from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType from databricks.sql.exc import RequestError +from databricks.sql.common import ssl_context_cache from databricks.sql.common.http import HttpMethod from databricks.sql.common.http_utils import ( detect_and_parse_proxy, @@ -94,31 +95,37 @@ def __init__(self, client_context): self._proxy_auth = None self._setup_pool_managers() + @staticmethod + def _create_ssl_context(ssl_options) -> ssl.SSLContext: + ssl_context = ssl.create_default_context() + + # Configure SSL verification + if not ssl_options.tls_verify: + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + elif not ssl_options.tls_verify_hostname: + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_REQUIRED + + # Load custom CA file if specified + if ssl_options.tls_trusted_ca_file: + ssl_context.load_verify_locations(ssl_options.tls_trusted_ca_file) + + # Load a separate cert/key pair or a combined cert+key PEM. The shared + # helper also reports missing/empty paths before opaque stdlib SSL errors. + ssl_options.load_client_cert_chain(ssl_context) + return ssl_context + def _setup_pool_managers(self): """Set up both direct and proxy pool managers for per-request proxy decisions.""" # SSL context setup ssl_context = None - if self.config.ssl_options: - ssl_context = ssl.create_default_context() - - # Configure SSL verification - if not self.config.ssl_options.tls_verify: - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - elif not self.config.ssl_options.tls_verify_hostname: - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_REQUIRED - - # Load custom CA file if specified - if self.config.ssl_options.tls_trusted_ca_file: - ssl_context.load_verify_locations( - self.config.ssl_options.tls_trusted_ca_file - ) - - # Load a separate cert/key pair or a combined cert+key PEM. The shared - # helper also reports missing/empty paths before opaque stdlib SSL errors. - self.config.ssl_options.load_client_cert_chain(ssl_context) + ssl_options = self.config.ssl_options + if ssl_options: + ssl_context = ssl_context_cache.get_ssl_context( + "unified", ssl_options, lambda: self._create_ssl_context(ssl_options) + ) # Create retry policy self._retry_policy = DatabricksRetryPolicy( diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 897d8774c..d2c1f6e5e 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,6 +2,16 @@ import pytest +from databricks.sql.common import ssl_context_cache + + +@pytest.fixture(autouse=True) +def clear_ssl_context_cache(): + # Tests patch the stdlib/urllib3 context builders; never serve another test's context. + ssl_context_cache.clear() + yield + ssl_context_cache.clear() + @pytest.fixture(autouse=True) def session_feature_flags(): diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py new file mode 100644 index 000000000..735bc84eb --- /dev/null +++ b/tests/unit/test_ssl_context_cache.py @@ -0,0 +1,243 @@ +import os +import ssl +import threading +import time +from unittest.mock import MagicMock, Mock, patch + +import pytest +from urllib3.util.ssl_ import create_urllib3_context + +import databricks.sql.backend.thrift_backend # noqa: F401 (loads thrift's THttpClient) +from databricks.sql.auth.common import ClientContext +from databricks.sql.auth.thrift_http_client import THttpClient +from databricks.sql.common import ssl_context_cache +from databricks.sql.common.unified_http_client import UnifiedHttpClient +from databricks.sql.types import SSLOptions + + +def counting_builder(): + built = [] + + def build(): + context = Mock(spec=ssl.SSLContext) + built.append(context) + return context + + return build, built + + +class TestGetSslContext: + def test_identical_options_share_one_context(self): + build, built = counting_builder() + + first = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + second = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + + assert first is second + assert len(built) == 1 + + @pytest.mark.parametrize( + "options", + [ + SSLOptions(tls_verify=False), + SSLOptions(tls_verify_hostname=False), + SSLOptions(tls_trusted_ca_file="/nonexistent/ca.pem"), + ], + ) + def test_different_options_do_not_share(self, options): + build, built = counting_builder() + + default = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + other = ssl_context_cache.get_ssl_context("k", options, build) + + assert default is not other + assert len(built) == 2 + + def test_kinds_do_not_share(self): + build, built = counting_builder() + + a = ssl_context_cache.get_ssl_context("a", SSLOptions(), build) + b = ssl_context_cache.get_ssl_context("b", SSLOptions(), build) + + assert a is not b + + @pytest.mark.parametrize( + "options", + [ + SSLOptions(tls_client_cert_file="cert.pem"), + SSLOptions(tls_client_cert_file="cert.pem", tls_client_cert_key_file="k"), + SSLOptions(tls_client_cert_key_file="key.pem"), + ], + ) + def test_client_certificates_are_never_cached(self, options): + build, built = counting_builder() + + first = ssl_context_cache.get_ssl_context("k", options, build) + second = ssl_context_cache.get_ssl_context("k", options, build) + + assert first is not second + assert len(built) == 2 + + def test_rewritten_ca_file_gets_a_new_context(self, tmp_path): + ca_file = tmp_path / "ca.pem" + ca_file.write_text("one") + options = SSLOptions(tls_trusted_ca_file=str(ca_file)) + build, built = counting_builder() + + before = ssl_context_cache.get_ssl_context("k", options, build) + ca_file.write_text("rotated") + os.utime(ca_file, ns=(0, 10**9)) + after = ssl_context_cache.get_ssl_context("k", options, build) + + assert before is not after + assert ssl_context_cache.get_ssl_context("k", options, build) is after + + def test_default_bundle_env_change_gets_a_new_context(self, tmp_path, monkeypatch): + bundle = tmp_path / "bundle.pem" + bundle.write_text("x") + build, built = counting_builder() + + default = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + monkeypatch.setenv("SSL_CERT_FILE", str(bundle)) + overridden = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + + assert default is not overridden + + def test_failed_build_is_not_cached(self): + build, built = counting_builder() + failing = Mock(side_effect=FileNotFoundError("ca.pem")) + + with pytest.raises(FileNotFoundError): + ssl_context_cache.get_ssl_context("k", SSLOptions(), failing) + context = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + + assert context is built[0] + + def test_concurrent_first_use_builds_once(self): + built = [] + + def slow_build(): + time.sleep(0.05) + context = Mock(spec=ssl.SSLContext) + built.append(context) + return context + + threads = 32 + barrier = threading.Barrier(threads) + results = [] + + def worker(): + barrier.wait() + results.append( + ssl_context_cache.get_ssl_context("k", SSLOptions(), slow_build) + ) + + workers = [threading.Thread(target=worker) for _ in range(threads)] + for t in workers: + t.start() + for t in workers: + t.join() + + assert len(built) == 1 + assert len(results) == threads + assert all(r is built[0] for r in results) + + def test_cache_is_bounded(self, monkeypatch): + monkeypatch.setattr(ssl_context_cache, "_MAX_ENTRIES", 2) + build, built = counting_builder() + + for kind in ("a", "b", "c"): + ssl_context_cache.get_ssl_context(kind, SSLOptions(), build) + + assert len(ssl_context_cache._contexts) == 2 + assert ssl_context_cache.get_ssl_context("a", SSLOptions(), build) is not ( + built[0] + ) + + +class TestThriftTransportContext: + @staticmethod + def pool_context(ssl_options): + client = THttpClient( + auth_provider=None, + uri_or_host="https://example.com/sql/1.0/warehouses/abc", + ssl_options=ssl_options, + ) + client.open() + pool = client._THttpClient__pool + assert pool.ca_certs is None and pool.cert_file is None + return pool.conn_kw["ssl_context"], pool + + @pytest.mark.parametrize( + "tls_verify, cert_reqs", [(True, ssl.CERT_REQUIRED), (False, ssl.CERT_NONE)] + ) + def test_context_matches_what_urllib3_builds(self, tls_verify, cert_reqs): + context, pool = self.pool_context(SSLOptions(tls_verify=tls_verify)) + expected = create_urllib3_context(cert_reqs=cert_reqs) + + assert pool.cert_reqs == cert_reqs + assert context.verify_mode == expected.verify_mode + assert context.check_hostname == expected.check_hostname + assert context.options == expected.options + assert context.verify_flags == expected.verify_flags + assert context.minimum_version == expected.minimum_version + assert context.hostname_checks_common_name == ( + expected.hostname_checks_common_name + ) + assert context.cert_store_stats()["x509_ca"] > 0 + + def test_connections_share_the_context(self): + first, _ = self.pool_context(SSLOptions()) + second, _ = self.pool_context(SSLOptions()) + unverified, _ = self.pool_context(SSLOptions(tls_verify=False)) + + assert first is second + assert unverified is not first + assert unverified.verify_mode == ssl.CERT_NONE + assert first.verify_mode == ssl.CERT_REQUIRED + + def test_custom_ca_file_replaces_default_certs(self, tmp_path): + ca_file = tmp_path / "ca.pem" + ca_file.write_text("pem") + context = MagicMock() + with patch( + "databricks.sql.auth.thrift_http_client.create_urllib3_context", + return_value=context, + ) as create: + pooled, _ = self.pool_context(SSLOptions(tls_trusted_ca_file=str(ca_file))) + + assert pooled is context + create.assert_called_once_with(cert_reqs=ssl.CERT_REQUIRED) + context.load_verify_locations.assert_called_once_with(str(ca_file)) + context.load_default_certs.assert_not_called() + + +class TestUnifiedHttpClientContext: + @staticmethod + def client(ssl_options): + return UnifiedHttpClient( + ClientContext(hostname="https://example.com", ssl_options=ssl_options) + ) + + @staticmethod + def pool_context(client): + return client._direct_pool_manager.connection_pool_kw["ssl_context"] + + def test_clients_share_the_context(self): + first = self.pool_context(self.client(SSLOptions())) + second = self.pool_context(self.client(SSLOptions())) + + assert first is second + assert first.verify_mode == ssl.CERT_REQUIRED and first.check_hostname + + def test_different_options_get_their_own_context(self): + verified = self.pool_context(self.client(SSLOptions())) + no_hostname = self.pool_context( + self.client(SSLOptions(tls_verify_hostname=False)) + ) + unverified = self.pool_context(self.client(SSLOptions(tls_verify=False))) + + assert len({id(verified), id(no_hostname), id(unverified)}) == 3 + assert not no_hostname.check_hostname + assert no_hostname.verify_mode == ssl.CERT_REQUIRED + assert unverified.verify_mode == ssl.CERT_NONE diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index f127ad8e2..9ea21b86d 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -776,12 +776,15 @@ def test_tls_cert_args_are_used_by_http_client( self.assertEqual(http_client.certfile, mock_ssl_options.tls_client_cert_file) self.assertEqual(http_client.keyfile, mock_ssl_options.tls_client_cert_key_file) self.assertIsNotNone(http_client.certfile) - mock_create_default_context.assert_called() + # The pool never used a context built here; urllib3 builds its own. + mock_create_default_context.assert_not_called() http_client.open() conn_pool = http_client._THttpClient__pool self.assertIsInstance(conn_pool, HTTPSConnectionPool) + # Client certificates keep urllib3's per-connection context. + self.assertNotIn("ssl_context", conn_pool.conn_kw) self.assertEqual(conn_pool.cert_reqs, CERT_REQUIRED) self.assertEqual(conn_pool.ca_certs, mock_ssl_options.tls_trusted_ca_file) self.assertEqual(conn_pool.cert_file, mock_ssl_options.tls_client_cert_file) From 165dbbd8b5feb2b9c4c8f231acfd570256a1605a Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:27:39 +0000 Subject: [PATCH 2/7] perf: build SSL contexts under a per-key lock Building a context while holding the cache's only lock made a slow CA file read delay cache hits and builds for unrelated settings. The shared lock now only guards the dicts; each key builds under its own lock. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/common/ssl_context_cache.py | 32 ++++++++++++++--- tests/unit/test_ssl_context_cache.py | 36 +++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py index 160f88e5e..ae6a2e277 100644 --- a/src/databricks/sql/common/ssl_context_cache.py +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -11,7 +11,7 @@ import ssl import threading from collections import OrderedDict -from typing import TYPE_CHECKING, Callable, Optional, Tuple +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple if TYPE_CHECKING: from databricks.sql.types import SSLOptions @@ -20,6 +20,7 @@ _lock = threading.Lock() _contexts: "OrderedDict[tuple, ssl.SSLContext]" = OrderedDict() +_build_locks: Dict[tuple, Any] = {} def _file_signature(path) -> Optional[Tuple]: @@ -66,25 +67,46 @@ def get_ssl_context( key = _cache_key(kind, ssl_options) with _lock: context = _contexts.get(key) - if context is None: + if context is not None: + _contexts.move_to_end(key) + return context + key_lock = _build_locks.setdefault(key, threading.Lock()) + + # Build outside the shared lock so a slow CA file only delays its own key. + with key_lock: + with _lock: + context = _contexts.get(key) + if context is not None: + return context + try: context = build() + except BaseException: + _forget_build_lock(key, key_lock) + raise + with _lock: _contexts[key] = context if len(_contexts) > _MAX_ENTRIES: _contexts.popitem(last=False) - else: - _contexts.move_to_end(key) + _forget_build_lock(key, key_lock) return context +def _forget_build_lock(key: tuple, key_lock: Any) -> None: + with _lock: + if _build_locks.get(key) is key_lock: + del _build_locks[key] + + def clear() -> None: with _lock: _contexts.clear() def _reset_after_fork() -> None: - # A fork during a build would leave the child's copy of the lock held. + # A fork during a build would leave the child's copies of the locks held. global _lock _lock = threading.Lock() + _build_locks.clear() if hasattr(os, "register_at_fork"): diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index 735bc84eb..184a52751 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -142,6 +142,42 @@ def worker(): assert len(results) == threads assert all(r is built[0] for r in results) + def test_slow_build_does_not_block_other_keys(self): + build, built = counting_builder() + cached = ssl_context_cache.get_ssl_context("cached", SSLOptions(), build) + release = threading.Event() + + def stuck_build(): + release.wait(10) + return Mock(spec=ssl.SSLContext) + + stuck = threading.Thread( + target=ssl_context_cache.get_ssl_context, + args=("stuck", SSLOptions(), stuck_build), + ) + stuck.start() + try: + results = {} + + def lookups(): + results["hit"] = ssl_context_cache.get_ssl_context( + "cached", SSLOptions(), build + ) + results["other"] = ssl_context_cache.get_ssl_context( + "other", SSLOptions(), build + ) + + other = threading.Thread(target=lookups) + other.start() + other.join(timeout=2) + + assert not other.is_alive() + assert results["hit"] is cached + assert results["other"] is built[1] + finally: + release.set() + stuck.join() + def test_cache_is_bounded(self, monkeypatch): monkeypatch.setattr(ssl_context_cache, "_MAX_ENTRIES", 2) build, built = counting_builder() From 379e1b792b8da4be4bd7cb96dd614339abdfed55 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 08:47:01 +0000 Subject: [PATCH 3/7] Document the shared-context invariant; make parity test host-independent urllib3 writes a pool's cert_reqs and ALPN protocols back onto its context per connection, so state that sharers must derive their TLS pool kwargs from the cached options. Compare the CA store against urllib3's own context instead of assuming the host has a CA bundle file. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/common/ssl_context_cache.py | 7 +++++-- tests/unit/test_ssl_context_cache.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py index ae6a2e277..8d6775255 100644 --- a/src/databricks/sql/common/ssl_context_cache.py +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -58,8 +58,11 @@ def get_ssl_context( ) -> ssl.SSLContext: """Return the shared context for ``kind`` and ``ssl_options``, building it once. - ``kind`` separates builders that configure contexts differently. Callers - must not modify the returned context. + ``kind`` separates builders that configure contexts differently. urllib3 + writes a pool's ``cert_reqs`` and ALPN protocols back onto its context on + every connection, so every pool sharing a context must use TLS pool kwargs + derived only from ``ssl_options`` (and no ``assert_hostname``). Do not + otherwise modify the returned context. """ if not is_cacheable(ssl_options): return build() diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index 184a52751..ccdbfe76c 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -210,6 +210,7 @@ def pool_context(ssl_options): def test_context_matches_what_urllib3_builds(self, tls_verify, cert_reqs): context, pool = self.pool_context(SSLOptions(tls_verify=tls_verify)) expected = create_urllib3_context(cert_reqs=cert_reqs) + expected.load_default_certs() assert pool.cert_reqs == cert_reqs assert context.verify_mode == expected.verify_mode @@ -220,7 +221,7 @@ def test_context_matches_what_urllib3_builds(self, tls_verify, cert_reqs): assert context.hostname_checks_common_name == ( expected.hostname_checks_common_name ) - assert context.cert_store_stats()["x509_ca"] > 0 + assert context.cert_store_stats() == expected.cert_store_stats() def test_connections_share_the_context(self): first, _ = self.pool_context(SSLOptions()) From b8d5fb0db4f6e0c3f2f6571219f310f7c89450d3 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 09:06:38 +0000 Subject: [PATCH 4/7] Rebuild cached SSL contexts after 10 minutes load_default_certs() also reads trust stores no file signature covers, such as the Windows certificate store, so expire entries to keep picking up system CA changes without a restart. Also cover the proxied Thrift pool in tests. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/common/ssl_context_cache.py | 25 ++++++++++++---- tests/unit/test_ssl_context_cache.py | 29 +++++++++++++++++++ 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py index 8d6775255..f40248ca7 100644 --- a/src/databricks/sql/common/ssl_context_cache.py +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -4,12 +4,15 @@ not scale across threads, so building one per connection caps how fast a process can open connections. A cached context is reused only by callers with identical TLS options and unchanged CA files. Configurations with a client -certificate are never cached. +certificate are never cached. Contexts are rebuilt after ``_TTL_SECONDS`` so +trust-store changes that no file signature covers (e.g. the Windows +certificate store) still apply without a restart. """ import os import ssl import threading +import time from collections import OrderedDict from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple @@ -17,9 +20,10 @@ from databricks.sql.types import SSLOptions _MAX_ENTRIES = 32 +_TTL_SECONDS = 600 _lock = threading.Lock() -_contexts: "OrderedDict[tuple, ssl.SSLContext]" = OrderedDict() +_contexts: "OrderedDict[tuple, Tuple[ssl.SSLContext, float]]" = OrderedDict() _build_locks: Dict[tuple, Any] = {} @@ -47,6 +51,15 @@ def _cache_key(kind: str, ssl_options: "SSLOptions") -> tuple: ) +def _fresh_context(key: tuple) -> Optional[ssl.SSLContext]: + # Caller holds _lock. + entry = _contexts.get(key) + if entry is None or time.monotonic() - entry[1] > _TTL_SECONDS: + return None + _contexts.move_to_end(key) + return entry[0] + + def is_cacheable(ssl_options: "SSLOptions") -> bool: return not ( ssl_options.tls_client_cert_file or ssl_options.tls_client_cert_key_file @@ -69,16 +82,15 @@ def get_ssl_context( key = _cache_key(kind, ssl_options) with _lock: - context = _contexts.get(key) + context = _fresh_context(key) if context is not None: - _contexts.move_to_end(key) return context key_lock = _build_locks.setdefault(key, threading.Lock()) # Build outside the shared lock so a slow CA file only delays its own key. with key_lock: with _lock: - context = _contexts.get(key) + context = _fresh_context(key) if context is not None: return context try: @@ -87,7 +99,8 @@ def get_ssl_context( _forget_build_lock(key, key_lock) raise with _lock: - _contexts[key] = context + _contexts[key] = (context, time.monotonic()) + _contexts.move_to_end(key) if len(_contexts) > _MAX_ENTRIES: _contexts.popitem(last=False) _forget_build_lock(key, key_lock) diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index ccdbfe76c..d5347c3d7 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -2,6 +2,7 @@ import ssl import threading import time +from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch import pytest @@ -178,6 +179,23 @@ def lookups(): release.set() stuck.join() + def test_expired_context_is_rebuilt(self, monkeypatch): + now = [1000.0] + monkeypatch.setattr( + ssl_context_cache, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + build, built = counting_builder() + + first = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + now[0] += ssl_context_cache._TTL_SECONDS + assert ssl_context_cache.get_ssl_context("k", SSLOptions(), build) is first + now[0] += 1 + rebuilt = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + + assert rebuilt is not first + assert ssl_context_cache.get_ssl_context("k", SSLOptions(), build) is rebuilt + assert len(built) == 2 + def test_cache_is_bounded(self, monkeypatch): monkeypatch.setattr(ssl_context_cache, "_MAX_ENTRIES", 2) build, built = counting_builder() @@ -233,6 +251,17 @@ def test_connections_share_the_context(self): assert unverified.verify_mode == ssl.CERT_NONE assert first.verify_mode == ssl.CERT_REQUIRED + def test_proxied_pool_shares_the_context(self): + direct, _ = self.pool_context(SSLOptions()) + with patch( + "databricks.sql.auth.thrift_http_client.detect_and_parse_proxy", + return_value=("http://proxy.example.com:3128", None), + ): + proxied, pool = self.pool_context(SSLOptions()) + + assert pool.proxy is not None and pool.proxy.host == "proxy.example.com" + assert proxied is direct + def test_custom_ca_file_replaces_default_certs(self, tmp_path): ca_file = tmp_path / "ca.pem" ca_file.write_text("pem") From c1fd8c6b1565e6455d223d9b1bf079c42dd0f6f5 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 09:26:36 +0000 Subject: [PATCH 5/7] Keep THttpClient's fail-fast check for client certificates The removed context build also validated the mTLS identity when the transport was constructed. Keep it for client-certificate configs only, which are not cached, so a key without a certificate still raises instead of silently using one-way TLS. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- src/databricks/sql/auth/thrift_http_client.py | 3 +++ tests/unit/test_ssl_context_cache.py | 19 +++++++++++++++++++ tests/unit/test_thrift_backend.py | 3 +-- 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 5236b9770..f6cbcbf37 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -73,6 +73,9 @@ def __init__( if self._ssl_options is not None: self.certfile = self._ssl_options.tls_client_cert_file self.keyfile = self._ssl_options.tls_client_cert_key_file + if not ssl_context_cache.is_cacheable(self._ssl_options): + # Fail fast on a bad client identity; urllib3 loads it lazily. + self._ssl_options.create_ssl_context() self.port = parsed.port self.host = parsed.hostname self.path = parsed.path diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index d5347c3d7..a6abae60d 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -13,6 +13,7 @@ from databricks.sql.auth.thrift_http_client import THttpClient from databricks.sql.common import ssl_context_cache from databricks.sql.common.unified_http_client import UnifiedHttpClient +from databricks.sql.exc import ProgrammingError from databricks.sql.types import SSLOptions @@ -251,6 +252,24 @@ def test_connections_share_the_context(self): assert unverified.verify_mode == ssl.CERT_NONE assert first.verify_mode == ssl.CERT_REQUIRED + def test_invalid_client_identity_fails_at_construction(self): + with pytest.raises(ProgrammingError, match="requires tls_client_cert_file"): + THttpClient( + auth_provider=None, + uri_or_host="https://example.com/sql/1.0/warehouses/abc", + ssl_options=SSLOptions(tls_client_cert_key_file="key.pem"), + ) + + def test_construction_without_client_identity_loads_no_context(self): + with patch("databricks.sql.types.create_default_context") as create: + THttpClient( + auth_provider=None, + uri_or_host="https://example.com/sql/1.0/warehouses/abc", + ssl_options=SSLOptions(), + ) + + create.assert_not_called() + def test_proxied_pool_shares_the_context(self): direct, _ = self.pool_context(SSLOptions()) with patch( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 9ea21b86d..39884b8ce 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -776,8 +776,7 @@ def test_tls_cert_args_are_used_by_http_client( self.assertEqual(http_client.certfile, mock_ssl_options.tls_client_cert_file) self.assertEqual(http_client.keyfile, mock_ssl_options.tls_client_cert_key_file) self.assertIsNotNone(http_client.certfile) - # The pool never used a context built here; urllib3 builds its own. - mock_create_default_context.assert_not_called() + mock_create_default_context.assert_called() http_client.open() From d4f71e8f9a893d0e72f89bd9bba62a7b195cadb2 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 09:46:28 +0000 Subject: [PATCH 6/7] Include ctime in the CA file signature An in-place rewrite that keeps the size and mtime (cp -p, rsync -t) left the key unchanged until the TTL expired; ctime cannot be preserved that way. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/common/ssl_context_cache.py | 3 ++- tests/unit/test_ssl_context_cache.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py index f40248ca7..06d56f897 100644 --- a/src/databricks/sql/common/ssl_context_cache.py +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -35,7 +35,8 @@ def _file_signature(path) -> Optional[Tuple]: st = os.stat(path) except (OSError, TypeError, ValueError): return (path, None) - return (path, st.st_ino, st.st_size, st.st_mtime_ns) + # ctime too: in-place rewrites that preserve mtime (cp -p, rsync -t) still change it. + return (path, st.st_ino, st.st_size, st.st_mtime_ns, st.st_ctime_ns) def _cache_key(kind: str, ssl_options: "SSLOptions") -> tuple: diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index a6abae60d..a3bea0261 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -1,5 +1,6 @@ import os import ssl +import sys import threading import time from types import SimpleNamespace @@ -94,6 +95,29 @@ def test_rewritten_ca_file_gets_a_new_context(self, tmp_path): assert before is not after assert ssl_context_cache.get_ssl_context("k", options, build) is after + @pytest.mark.skipif( + sys.platform == "win32", reason="st_ctime is the creation time on Windows" + ) + def test_in_place_rewrite_keeping_mtime_gets_a_new_context(self, tmp_path): + ca_file = tmp_path / "ca.pem" + ca_file.write_text("old") + original = ca_file.stat() + options = SSLOptions(tls_trusted_ca_file=str(ca_file)) + build, built = counting_builder() + before = ssl_context_cache.get_ssl_context("k", options, build) + + time.sleep(0.05) # filesystem timestamps can be jiffy-granular + ca_file.write_text("new") # same inode and size, like `cp -p` + os.utime(ca_file, ns=(original.st_atime_ns, original.st_mtime_ns)) + rewritten = ca_file.stat() + + assert (rewritten.st_ino, rewritten.st_size, rewritten.st_mtime_ns) == ( + original.st_ino, + original.st_size, + original.st_mtime_ns, + ) + assert ssl_context_cache.get_ssl_context("k", options, build) is not before + def test_default_bundle_env_change_gets_a_new_context(self, tmp_path, monkeypatch): bundle = tmp_path / "bundle.pem" bundle.write_text("x") From 78f1e1425aa4d6b76d400836895bfc8e90b425c5 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 12 Nov 2026 10:12:46 +0000 Subject: [PATCH 7/7] Stop urllib3 from rewriting shared SSL contexts per connection urllib3 calls set_alpn_protocols() and sets verify_mode on its context for every connection. On a shared context, replacing the ALPN buffer races with SSL_new() in other threads (which runs without the GIL) and can corrupt the ClientHello: 1-4 corrupted ALPN lists per ~500k handshake starts in a 32-thread stress test. Cached contexts now ignore same-value writes, with ALPN set once before the context is shared. Contexts that are not plain ssl.SSLContext objects are not cached. Signed-off-by: Vu Anh Phung Co-authored-by: Isaac --- .../sql/common/ssl_context_cache.py | 49 +++++++++++- tests/unit/test_ssl_context_cache.py | 75 ++++++++++++++++++- 2 files changed, 118 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/common/ssl_context_cache.py b/src/databricks/sql/common/ssl_context_cache.py index 06d56f897..b565bde73 100644 --- a/src/databricks/sql/common/ssl_context_cache.py +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -16,6 +16,8 @@ from collections import OrderedDict from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple +from urllib3.util.ssl_ import ALPN_PROTOCOLS + if TYPE_CHECKING: from databricks.sql.types import SSLOptions @@ -26,6 +28,41 @@ _contexts: "OrderedDict[tuple, Tuple[ssl.SSLContext, float]]" = OrderedDict() _build_locks: Dict[tuple, Any] = {} +_BASE_CONTEXT_CLASS = ssl.SSLContext + + +class _SharedSSLContext(_BASE_CONTEXT_CLASS): # type: ignore[valid-type,misc] + """A cached context that ignores same-value writes. + + urllib3 re-applies ALPN and verify_mode to its context on every connection. + SSL_new() copies the context's ALPN buffer without the GIL, so replacing + that buffer while another thread connects can corrupt the ClientHello. + """ + + def set_alpn_protocols(self, alpn_protocols): + protocols = list(alpn_protocols) + if protocols != getattr(self, "_shared_alpn_protocols", None): + super().set_alpn_protocols(protocols) + self._shared_alpn_protocols = protocols + + @property + def verify_mode(self): + return super().verify_mode + + @verify_mode.setter + def verify_mode(self, value): + if super().verify_mode != value: + _BASE_CONTEXT_CLASS.verify_mode.__set__(self, value) + + +def _make_shareable(context: ssl.SSLContext) -> bool: + # Only plain stdlib contexts are shared; others (e.g. pyOpenSSL) stay per caller. + if type(context) is not _BASE_CONTEXT_CLASS: + return False + context.__class__ = _SharedSSLContext + context.set_alpn_protocols(ALPN_PROTOCOLS) + return True + def _file_signature(path) -> Optional[Tuple]: # Rewriting or replacing a CA file yields a new key, so rotation is picked up. @@ -74,9 +111,11 @@ def get_ssl_context( ``kind`` separates builders that configure contexts differently. urllib3 writes a pool's ``cert_reqs`` and ALPN protocols back onto its context on - every connection, so every pool sharing a context must use TLS pool kwargs - derived only from ``ssl_options`` (and no ``assert_hostname``). Do not - otherwise modify the returned context. + every connection; the shared context skips those writes only when nothing + changes, so every pool sharing a context must use TLS pool kwargs derived + only from ``ssl_options`` (and no ``assert_hostname``). Do not otherwise + modify the returned context. Contexts other than plain ``ssl.SSLContext`` + are returned without caching. """ if not is_cacheable(ssl_options): return build() @@ -96,9 +135,13 @@ def get_ssl_context( return context try: context = build() + shareable = _make_shareable(context) except BaseException: _forget_build_lock(key, key_lock) raise + if not shareable: + _forget_build_lock(key, key_lock) + return context with _lock: _contexts[key] = (context, time.monotonic()) _contexts.move_to_end(key) diff --git a/tests/unit/test_ssl_context_cache.py b/tests/unit/test_ssl_context_cache.py index a3bea0261..fcab6ffcb 100644 --- a/tests/unit/test_ssl_context_cache.py +++ b/tests/unit/test_ssl_context_cache.py @@ -18,11 +18,16 @@ from databricks.sql.types import SSLOptions +def bare_context(): + # Real but cheap: no CA bundle is loaded. + return ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + + def counting_builder(): built = [] def build(): - context = Mock(spec=ssl.SSLContext) + context = bare_context() built.append(context) return context @@ -144,7 +149,7 @@ def test_concurrent_first_use_builds_once(self): def slow_build(): time.sleep(0.05) - context = Mock(spec=ssl.SSLContext) + context = bare_context() built.append(context) return context @@ -175,7 +180,7 @@ def test_slow_build_does_not_block_other_keys(self): def stuck_build(): release.wait(10) - return Mock(spec=ssl.SSLContext) + return bare_context() stuck = threading.Thread( target=ssl_context_cache.get_ssl_context, @@ -233,6 +238,70 @@ def test_cache_is_bounded(self, monkeypatch): built[0] ) + def test_non_stdlib_contexts_are_not_cached(self): + build = Mock(side_effect=lambda: Mock(spec=ssl.SSLContext)) + + first = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + second = ssl_context_cache.get_ssl_context("k", SSLOptions(), build) + + assert first is not second + assert not ssl_context_cache._contexts + + def test_shared_context_skips_same_value_writes(self): + context = ssl_context_cache.get_ssl_context("k", SSLOptions(), bare_context) + assert isinstance(context, ssl.SSLContext) + + with patch.object(ssl.SSLContext, "_set_alpn_protocols") as set_alpn: + context.set_alpn_protocols(["http/1.1"]) # what urllib3 sends + set_alpn.assert_not_called() + context.set_alpn_protocols(["h2"]) + set_alpn.assert_called_once() + + real = ssl.SSLContext.verify_mode + writes = [] + + def spy_set(self, value): + writes.append(value) + real.fset(self, value) + + context.check_hostname = False + with patch.object(ssl.SSLContext, "verify_mode", property(real.fget, spy_set)): + context.verify_mode = ssl.CERT_REQUIRED + assert writes == [] + context.verify_mode = ssl.CERT_NONE + assert writes == [ssl.CERT_NONE] + assert context.verify_mode == ssl.CERT_NONE + + def test_reset_after_fork_releases_locks(self): + held = threading.Event() + release = threading.Event() + + def hold_locks(): + with ssl_context_cache._lock: + held.set() + release.wait(10) + + holder = threading.Thread(target=hold_locks) + holder.start() + held.wait(5) + stuck_build_lock = threading.Lock() + stuck_build_lock.acquire() + key = ssl_context_cache._cache_key("k", SSLOptions()) + ssl_context_cache._build_locks[key] = stuck_build_lock + try: + ssl_context_cache._reset_after_fork() # as in a forked child + + done = threading.Thread( + target=ssl_context_cache.get_ssl_context, + args=("k", SSLOptions(), bare_context), + ) + done.start() + done.join(timeout=2) + assert not done.is_alive() + finally: + release.set() + holder.join() + class TestThriftTransportContext: @staticmethod