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..f6cbcbf37 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,11 @@ 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() + 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 @@ -122,19 +135,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..b565bde73 --- /dev/null +++ b/src/databricks/sql/common/ssl_context_cache.py @@ -0,0 +1,173 @@ +"""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. 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 + +from urllib3.util.ssl_ import ALPN_PROTOCOLS + +if TYPE_CHECKING: + from databricks.sql.types import SSLOptions + +_MAX_ENTRIES = 32 +_TTL_SECONDS = 600 + +_lock = threading.Lock() +_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. + if not path: + return None + try: + st = os.stat(path) + except (OSError, TypeError, ValueError): + return (path, None) + # 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: + 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 _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 + ) + + +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. urllib3 + writes a pool's ``cert_reqs`` and ALPN protocols back onto its context on + 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() + + key = _cache_key(kind, ssl_options) + with _lock: + context = _fresh_context(key) + if context is not None: + 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 = _fresh_context(key) + if context is not None: + 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) + if len(_contexts) > _MAX_ENTRIES: + _contexts.popitem(last=False) + _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 copies of the locks held. + global _lock + _lock = threading.Lock() + _build_locks.clear() + + +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..fcab6ffcb --- /dev/null +++ b/tests/unit/test_ssl_context_cache.py @@ -0,0 +1,421 @@ +import os +import ssl +import sys +import threading +import time +from types import SimpleNamespace +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.exc import ProgrammingError +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 = bare_context() + 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 + + @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") + 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 = bare_context() + 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_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 bare_context() + + 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_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() + + 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] + ) + + 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 + 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) + expected.load_default_certs() + + 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() == expected.cert_store_stats() + + 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_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( + "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") + 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..39884b8ce 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -782,6 +782,8 @@ def test_tls_cert_args_are_used_by_http_client( 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)