SNOW-2395236: renamed to use_requests_session to use_session - to improve extensibility of code. Made SessionPool generic.

Author: sfc-gh-fpawlowski

ci/pre-commit/check_no_native_http.py

@@ -869,7 +869,7 @@ def _check_requests_call(
                 node.lineno,
                 node.col_offset,
                 ViolationType.REQUESTS_SESSION,
-                "Direct use of requests.Session() is forbidden, use SessionManager.use_requests_session() instead",
+                "Direct use of requests.Session() is forbidden, use SessionManager.use_session() instead",
             )
         elif ModulePattern.is_http_method(func_name):
             return HTTPViolation(
@@ -1039,11 +1039,9 @@ def main():
             print()
             print("How to fix:")
             print("  - Replace requests.request() with SessionManager.request()")
+            print("  - Replace requests.Session() with SessionManager.use_session()")
             print(
-                "  - Replace requests.Session() with SessionManager.use_requests_session()"
-            )
-            print(
-                "  - Replace urllib3.PoolManager/ProxyManager() with session from session_manager.use_requests_session()"
+                "  - Replace urllib3.PoolManager/ProxyManager() with session from session_manager.use_session()"
             )
             print("  - Replace direct HTTP method imports with SessionManager usage")
             print("  - Use SessionManager for all HTTP operations")

src/snowflake/connector/network.py

@@ -851,7 +851,7 @@ def add_retry_params(self, full_url: str) -> str:
         include_retry_reason = self._connection._enable_retry_reason_in_query_response
         include_retry_params = kwargs.pop("_include_retry_params", False)
 
-        with self.use_requests_session(full_url) as session:
+        with self.use_session(full_url) as session:
             retry_ctx = RetryCtx(
                 _include_retry_params=include_retry_params,
                 _include_retry_reason=include_retry_reason,
@@ -1213,5 +1213,5 @@ def _request_exec(
         except Exception as err:
             raise err
 
-    def use_requests_session(self, url=None) -> Generator[Session, Any, None]:
-        return self.session_manager.use_requests_session(url)
+    def use_session(self, url=None) -> Generator[Session, Any, None]:
+        return self.session_manager.use_session(url)

src/snowflake/connector/ocsp_snowflake.py

@@ -552,7 +552,7 @@ def _download_ocsp_response_cache(ocsp, url, do_retry: bool = True) -> bool:
             session_manager = get_current_session_manager(
                 use_pooling=False
             ) or SessionManager(use_pooling=False)
-            with session_manager.use_requests_session() as session:
+            with session_manager.use_session() as session:
                 max_retry = SnowflakeOCSP.OCSP_CACHE_SERVER_MAX_RETRY if do_retry else 1
                 sleep_time = 1
                 backoff = exponential_backoff()()
@@ -1648,7 +1648,7 @@ def _fetch_ocsp_response(
             if context_session_manager is not None
             else SessionManager(use_pooling=False)
         )
-        with session_manager.use_requests_session() as session:
+        with session_manager.use_session() as session:
             max_retry = sf_max_retry if do_retry else 1
             sleep_time = 1
             backoff = exponential_backoff()()

src/snowflake/connector/result_batch.py

@@ -360,14 +360,14 @@ def _download(
                         and connection.rest.session_manager is not None
                     ):
                         # If connection was explicitly passed and not closed yet - we can reuse SessionManager with session pooling
-                        with connection.rest.use_requests_session() as session:
+                        with connection.rest.use_session() as session:
                             logger.debug(
                                 f"downloading result batch id: {self.id} with existing session {session}"
                             )
                             response = session.request("get", **request_data)
                     elif self._session_manager is not None:
                         # If connection is not accessible or was already closed, but cursors are now used to fetch the data - we will only reuse the http setup (through cloned SessionManager without session pooling)
-                        with self._session_manager.use_requests_session() as session:
+                        with self._session_manager.use_session() as session:
                             response = session.request("get", **request_data)
                     else:
                         # If there was no session manager cloned, then we are using a default Session Manager setup, since it is very unlikely to enter this part outside of testing

src/snowflake/connector/session_manager.py

@@ -7,7 +7,7 @@
 import itertools
 import logging
 from dataclasses import dataclass, field, replace
-from typing import TYPE_CHECKING, Any, Callable, Generator, Mapping
+from typing import TYPE_CHECKING, Any, Callable, Generator, Generic, Mapping, TypeVar
 
 from .compat import urlparse
 from .proxy import get_proxy_url
@@ -26,6 +26,9 @@
 logger = logging.getLogger(__name__)
 REQUESTS_RETRY = 1  # requests library builtin retry
 
+# Generic type for session objects (requests.Session, aiohttp.ClientSession, etc.) - no specific interface is required
+SessionT = TypeVar("SessionT")
+
 
 def _propagate_session_manager_to_ocsp(generator_func):
     """Decorator: push self into ssl_wrap_socket ContextVar for OCSP duration.
@@ -122,23 +125,29 @@ def __call__(self, *args, **kwargs) -> ProxySupportAdapter:
 
 
 @dataclass(frozen=True)
-class HttpConfig:
+class BaseHttpConfig:
     """Immutable HTTP configuration shared by SessionManager instances."""
 
-    adapter_factory: Callable[..., HTTPAdapter] = field(
-        default_factory=ProxySupportAdapterFactory
-    )
     use_pooling: bool = True
     max_retries: int | Retry | None = REQUESTS_RETRY
     proxy_host: str | None = None
     proxy_port: str | None = None
     proxy_user: str | None = None
     proxy_password: str | None = None
 
-    def copy_with(self, **overrides: Any) -> HttpConfig:
-        """Return a new HttpConfig with overrides applied."""
+    def copy_with(self, **overrides: Any) -> BaseHttpConfig:
+        """Return a new config with overrides applied."""
         return replace(self, **overrides)
 
+
+@dataclass(frozen=True)
+class HttpConfig(BaseHttpConfig):
+    """HTTP configuration specific to requests library."""
+
+    adapter_factory: Callable[..., HTTPAdapter] = field(
+        default_factory=ProxySupportAdapterFactory
+    )
+
     def get_adapter(self, **override_adapter_factory_kwargs) -> HTTPAdapter:
         # We pass here only chosen attributes as kwargs to make the arguments received by the factory as compliant with the HttpAdapter constructor interface as possible.
         # We could consider passing the whole HttpConfig as kwarg to the factory if necessary in the future.
@@ -156,9 +165,9 @@ def get_adapter(self, **override_adapter_factory_kwargs) -> HTTPAdapter:
         return self.adapter_factory(**self_kwargs_for_adapter_factory)
 
 
-class SessionPool:
+class SessionPool(Generic[SessionT]):
     """
-    Component responsible for storing and reusing established instances of requests.Session class.
+    Component responsible for storing and reusing established session instances.
 
     This approach is especially useful in scenarios where multiple requests would have to be sent
     to the same host in short period of time. Instead of repeatedly establishing a new TCP connection
@@ -167,15 +176,17 @@ class SessionPool:
 
     Sessions are created using the factory method make_session of a passed instance of the
     SessionManager class.
+
+    Generic over SessionT to support different session types (requests.Session, aiohttp.ClientSession, etc.)
     """
 
     def __init__(self, manager: SessionManager) -> None:
         # A stack of the idle sessions
-        self._idle_sessions = []
-        self._active_sessions = set()
+        self._idle_sessions: list[SessionT] = []
+        self._active_sessions: set[SessionT] = set()
         self._manager = manager
 
-    def get_session(self) -> Session:
+    def get_session(self) -> SessionT:
         """Returns a session from the session pool or creates a new one."""
         try:
             session = self._idle_sessions.pop()
@@ -184,7 +195,7 @@ def get_session(self) -> Session:
         self._active_sessions.add(session)
         return session
 
-    def return_session(self, session: Session) -> None:
+    def return_session(self, session: SessionT) -> None:
         """Places an active session back into the idle session stack."""
         try:
             self._active_sessions.remove(session)
@@ -249,11 +260,11 @@ class _RequestVerbsUsingSessionMixin(abc.ABC):
     """
     Mixin that provides HTTP methods (get, post, put, etc.) mirroring requests.Session, maintaining their default argument behavior (e.g., HEAD uses allow_redirects=False).
     These wrappers manage the SessionManager's use of pooled/non-pooled sessions and delegate the actual request to the corresponding session.<verb>() method.
-    The subclass must implement use_requests_session to yield a *requests.Session* instance.
+    The subclass must implement use_session to yield a *requests.Session* instance.
     """
 
     @abc.abstractmethod
-    def use_requests_session(self, url: str, use_pooling: bool) -> Session: ...
+    def use_session(self, url: str, use_pooling: bool) -> Session: ...
 
     def get(
         self,
@@ -264,7 +275,7 @@ def get(
         use_pooling: bool | None = None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.get(url, headers=headers, timeout=timeout, **kwargs)
 
     def options(
@@ -276,7 +287,7 @@ def options(
         use_pooling: bool | None = None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.options(url, headers=headers, timeout=timeout, **kwargs)
 
     def head(
@@ -288,7 +299,7 @@ def head(
         use_pooling: bool | None = None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.head(url, headers=headers, timeout=timeout, **kwargs)
 
     def post(
@@ -302,7 +313,7 @@ def post(
         json=None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.post(
                 url,
                 headers=headers,
@@ -322,7 +333,7 @@ def put(
         data=None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.put(
                 url, headers=headers, timeout=timeout, data=data, **kwargs
             )
@@ -337,7 +348,7 @@ def patch(
         data=None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.patch(
                 url, headers=headers, timeout=timeout, data=data, **kwargs
             )
@@ -351,7 +362,7 @@ def delete(
         use_pooling: bool | None = None,
         **kwargs,
     ):
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.delete(url, headers=headers, timeout=timeout, **kwargs)
 
 
@@ -466,7 +477,7 @@ def make_session(self) -> Session:
 
     @contextlib.contextmanager
     @_propagate_session_manager_to_ocsp
-    def use_requests_session(
+    def use_session(
         self, url: str | bytes | None = None, use_pooling: bool | None = None
     ) -> Generator[Session, Any, None]:
         use_pooling = use_pooling if use_pooling is not None else self.use_pooling
@@ -500,7 +511,7 @@ def request(
         This wraps :pymeth:`use_session` so callers don’t have to manage the
         context manager themselves.
         """
-        with self.use_requests_session(url, use_pooling) as session:
+        with self.use_session(url, use_pooling) as session:
             return session.request(
                 method=method.upper(),
                 url=url,

src/snowflake/connector/storage_client.py

@@ -289,7 +289,7 @@ def _send_request_with_retry(
             rest_kwargs["timeout"] = (REQUEST_CONNECTION_TIMEOUT, REQUEST_READ_TIMEOUT)
             try:
                 if conn:
-                    with conn.rest.use_requests_session(url=url) as session:
+                    with conn.rest.use_session(url=url) as session:
                         logger.debug(f"storage client request with session {session}")
                         response = session.request(verb, url, **rest_kwargs)
                 else:

src/snowflake/connector/telemetry_oob.py

@@ -482,7 +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(SNOW-2259522): 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).
+            # TODO(SNOW-2259522): 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_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",

test/integ/pandas_it/test_arrow_pandas.py

@@ -1429,8 +1429,8 @@ def test_sessions_used(conn_cnx, fetch_fn_name, pass_connection):
 
             # check that sessions are used when connection is supplied
             with mock.patch(
-                "snowflake.connector.network.SnowflakeRestful.use_requests_session",
-                side_effect=cnx._rest.use_requests_session,
+                "snowflake.connector.network.SnowflakeRestful.use_session",
+                side_effect=cnx._rest.use_session,
             ) as get_session_mock:
                 fetch_fn(connection=connection)
                 assert get_session_mock.call_count == (1 if pass_connection else 0)

test/integ/test_connection.py

@@ -1077,7 +1077,7 @@ def test_ocsp_and_rest_pool_isolation(conn_cnx, disable_request_pooling):
 
     assert rest_sm_1.sessions_map or disable_request_pooling
 
-    with rest_sm_1.use_requests_session("https://example.com"):
+    with rest_sm_1.use_session("https://example.com"):
         ocsp_sm_1 = get_current_session_manager(create_default_if_missing=False)
         assert ocsp_sm_1 is not rest_sm_1
         assert ocsp_sm_1.config == rest_sm_1.config
@@ -1096,7 +1096,7 @@ def test_ocsp_and_rest_pool_isolation(conn_cnx, disable_request_pooling):
     assert rest_sm_2.sessions_map or disable_request_pooling
     assert rest_sm_2 is not rest_sm_1
 
-    with rest_sm_2.use_requests_session("https://example.com"):
+    with rest_sm_2.use_session("https://example.com"):
         ocsp_sm_2 = get_current_session_manager(create_default_if_missing=False)
         assert ocsp_sm_2 is not rest_sm_2
         assert ocsp_sm_2.config == rest_sm_2.config

test/integ/test_cursor.py

@@ -1802,8 +1802,8 @@ def test_fetch_batches_with_sessions(conn_cnx):
             num_batches = len(cur.get_result_batches())
 
             with mock.patch(
-                "snowflake.connector.session_manager.SessionManager.use_requests_session",
-                side_effect=con._rest.session_manager.use_requests_session,
+                "snowflake.connector.session_manager.SessionManager.use_session",
+                side_effect=con._rest.session_manager.use_session,
             ) as get_session_mock:
                 result = cur.fetchall()
                 # all but one batch is downloaded using a session