SNOW-2203079: OCSP and Telemetry refactored

Author: sfc-gh-fpawlowski

src/snowflake/connector/connection.py

@@ -121,6 +121,9 @@
 )
 from .session_manager import SessionManager
 from .sqlstate import SQLSTATE_CONNECTION_NOT_EXISTS, SQLSTATE_FEATURE_NOT_SUPPORTED
+from .ssl_wrap_socket import (
+    set_current_session_manager as set_current_session_manager_for_ssl,
+)
 from .telemetry import TelemetryClient, TelemetryData, TelemetryField
 from .time_util import HeartBeatTimer, get_time_millis
 from .url_util import extract_top_level_domain_from_hostname
@@ -892,6 +895,15 @@ def connect(self, **kwargs) -> None:
         self._session_manager = SessionManager(
             use_pooling=(not self.disable_request_pooling),
         )
+        try:
+            # Set context var for OCSP downloads to use this manager. Can be removed once OCSP is deprecated.
+            self._ocsp_sm_context_var_old_token = set_current_session_manager_for_ssl(
+                self._session_manager
+            )
+        except Exception:
+            logger.debug(
+                "Could not set current SessionManager in ssl_wrap_socket", exc_info=True
+            )
 
         if self.enable_connection_diag:
             exceptions_dict = {}
@@ -932,6 +944,15 @@ def connect(self, **kwargs) -> None:
         else:
             self.__open_connection()
 
+        try:
+            # Reet context var for OCSP downloads to use this manager. Can be removed once OCSP is deprecated.
+            from .ssl_wrap_socket import reset_current_session_manager as _reset_sm
+
+            if hasattr(self, "_ocsp_sm_context_var_old_token"):
+                _reset_sm(self._ocsp_sm_context_var_old_token)
+        except Exception:
+            logger.debug("Failed to reset OCSP SessionManager context", exc_info=True)
+
     def close(self, retry: bool = True) -> None:
         """Closes the connection."""
         # unregister to dereference connection object as it's already closed after the execution

src/snowflake/connector/ocsp_snowflake.py

@@ -53,6 +53,8 @@
 )
 from snowflake.connector.errors import RevocationCheckError
 from snowflake.connector.network import PYTHON_CONNECTOR_USER_AGENT
+from snowflake.connector.session_manager import SessionManager
+from snowflake.connector.ssl_wrap_socket import get_current_session_manager
 
 from . import constants
 from .backoff_policies import exponential_backoff
@@ -546,7 +548,11 @@ def _download_ocsp_response_cache(ocsp, url, do_retry: bool = True) -> bool:
                 if sf_cache_server_url is not None:
                     url = sf_cache_server_url
 
-            with generic_requests.Session() as session:
+            # Obtain SessionManager from ssl_wrap_socket context var if available
+            session_manager = get_current_session_manager().clone(
+                use_pooling=False
+            ) or SessionManager(use_pooling=False)
+            with session_manager.use_requests_session() as session:
                 max_retry = SnowflakeOCSP.OCSP_CACHE_SERVER_MAX_RETRY if do_retry else 1
                 sleep_time = 1
                 backoff = exponential_backoff()()

src/snowflake/connector/ssl_wrap_socket.py

@@ -9,17 +9,23 @@
 # and added OCSP validator on the top.
 import logging
 import time
+import weakref
+from contextvars import ContextVar
 from functools import wraps
 from inspect import getfullargspec as get_args
 from socket import socket
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import certifi
 import OpenSSL.SSL
 
 from .constants import OCSPMode
 from .errorcode import ER_OCSP_RESPONSE_CERT_STATUS_REVOKED
 from .errors import OperationalError
+
+if TYPE_CHECKING:
+    from .session_manager import SessionManager
+
 from .vendored.urllib3 import connection as connection_
 from .vendored.urllib3.contrib.pyopenssl import PyOpenSSLContext, WrappedSocket
 from .vendored.urllib3.util import ssl_ as ssl_
@@ -35,6 +41,44 @@
 log = logging.getLogger(__name__)
 
 
+# Store a *weak* reference so that the context variable doesn’t prolong the
+# lifetime of the SessionManager. Once all owning connections are GC-ed the
+# weakref goes dead and OCSP will fall back to its local manager (but most likely won't be used ever again anyway).
+_CURRENT_SESSION_MANAGER: ContextVar[weakref.ref[SessionManager] | None] = ContextVar(
+    "_CURRENT_SESSION_MANAGER",
+    default=None,
+)
+
+
+def get_current_session_manager() -> SessionManager | None:
+    """Return the SessionManager associated with the current handshake, if any.
+
+    If the weak reference is dead or no manager was set, returns ``None``.
+    """
+    ref = _CURRENT_SESSION_MANAGER.get()
+    return ref() if ref is not None else None
+
+
+def set_current_session_manager(sm: SessionManager | None) -> Any:
+    """Set the SessionManager for the current execution context.
+
+    Called from SnowflakeConnection so that OCSP downloads
+    use the same proxy / header configuration as the initiating connection.
+
+    Alternative approach would be moving method inject_into_urllib3() inside connection initialization, but in case this delay (from module import time to connection initialization time) would cause some code to break we stayed with this approach, having in mind soon OCSP deprecation.
+    """
+    return _CURRENT_SESSION_MANAGER.set(weakref.ref(sm) if sm is not None else None)
+
+
+def reset_current_session_manager(token) -> None:
+    """Restore previous SessionManager context stored in *token* (from ContextVar.set)."""
+    try:
+        _CURRENT_SESSION_MANAGER.reset(token)
+    except Exception:
+        # ignore invalid token errors
+        pass
+
+
 def inject_into_urllib3() -> None:
     """Monkey-patch urllib3 with PyOpenSSL-backed SSL-support and OCSP."""
     log.debug("Injecting ssl_wrap_socket_with_ocsp")

src/snowflake/connector/telemetry_oob.py

@@ -482,6 +482,7 @@ def _upload_payload(self, payload) -> None:
                 # This logger guarantees the payload won't be masked. Testing purpose.
                 rt_plain_logger.debug(f"OOB telemetry data being sent is {payload}")
 
+            # TODO: Telemetry OOB is currently disabled. If Telemetry OOB is to be re-enabled, this HTTP call must be routed through the connection_argument.session_manager.use_requests_session(use_pooling) (so the SessionManager instance attached to the connection which initialization's fail most likely triggered this telemetry log). It would allow to pick up proxy configuration & custom headers (see tickets SNOW-694457 and SNOW-2203079).
             with requests.Session() as session:
                 headers = {
                     "Content-type": "application/json",