perf: share SSL contexts across connections to speed up connect()

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.

src/databricks/sql/auth/thrift_http_client.py

@@ -14,15 +14,27 @@
 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,
 )
 
 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(

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)

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(

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():