[async] Applied #2429 to async code - part 2:

Storage client
ResultBatch
WIF
replacing aiohttp with session manager and propagating it down partially - ocsp and session_manager file merge with _ssl_connector - analogically to ProxySupportAdapter

Author: sfc-gh-fpawlowski

.pre-commit-config.yaml

@@ -57,6 +57,7 @@ repos:
         exclude: |
             (?x)^(
                 src/snowflake/connector/session_manager\.py|
+                src/snowflake/connector/aio/_session_manager\.py|
                 src/snowflake/connector/vendored/.*
             )$
         args: [--show-fixes]

src/snowflake/connector/aio/_connection.py

@@ -200,7 +200,7 @@ async def __open_connection(self):
             protocol=self._protocol,
             inject_client_pause=self._inject_client_pause,
             connection=self,
-            session_manager=self._session_manager,
+            session_manager=self._session_manager,  # connection shares the session pool used for making Backend related requests
         )
         logger.debug("REST API object was created: %s:%s", self.host, self.port)
 
@@ -592,6 +592,7 @@ def _init_connection_parameters(
             PLATFORM,
         )
 
+        # Placeholder attributes; will be initialized in connect()
         self._http_config: AioHttpConfig | None = None
         self._session_manager: SessionManager | None = None
         self._rest = None

src/snowflake/connector/aio/_network.py

@@ -567,7 +567,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)
 
-        async with self._use_session(full_url) as session:
+        async with self.use_session(full_url) as session:
             retry_ctx = RetryCtx(
                 _include_retry_params=include_retry_params,
                 _include_retry_reason=include_retry_reason,

src/snowflake/connector/aio/_ocsp_snowflake.py

@@ -5,9 +5,8 @@
 import os
 import time
 from logging import getLogger
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-import aiohttp
 from aiohttp.client_proto import ResponseHandler
 from asn1crypto.ocsp import CertId
 from asn1crypto.x509 import Certificate
@@ -32,17 +31,22 @@
 from snowflake.connector.ocsp_snowflake import SnowflakeOCSP as SnowflakeOCSPSync
 from snowflake.connector.url_util import extract_top_level_domain_from_hostname
 
+if TYPE_CHECKING:
+    from snowflake.connector.aio._session_manager import SessionManager
+
 logger = getLogger(__name__)
 
 
 class OCSPServer(OCSPServerSync):
-    async def download_cache_from_server(self, ocsp):
+    async def download_cache_from_server(
+        self, ocsp, *, session_manager: SessionManager
+    ):
         if self.CACHE_SERVER_ENABLED:
             # if any of them is not cache, download the cache file from
             # OCSP response cache server.
             try:
                 retval = await OCSPServer._download_ocsp_response_cache(
-                    ocsp, self.CACHE_SERVER_URL
+                    ocsp, self.CACHE_SERVER_URL, session_manager=session_manager
                 )
                 if not retval:
                     raise RevocationCheckError(
@@ -69,7 +73,9 @@ async def download_cache_from_server(self, ocsp):
                 raise
 
     @staticmethod
-    async def _download_ocsp_response_cache(ocsp, url, do_retry: bool = True) -> bool:
+    async def _download_ocsp_response_cache(
+        ocsp, url, *, session_manager: SessionManager, do_retry: bool = True
+    ) -> bool:
         """Downloads OCSP response cache from the cache server."""
         headers = {HTTP_HEADER_USER_AGENT: PYTHON_CONNECTOR_USER_AGENT}
         sf_timeout = SnowflakeOCSP.OCSP_CACHE_SERVER_CONNECTION_TIMEOUT
@@ -88,7 +94,7 @@ async def _download_ocsp_response_cache(ocsp, url, do_retry: bool = True) -> boo
                 if sf_cache_server_url is not None:
                     url = sf_cache_server_url
 
-            async with aiohttp.ClientSession() as session:
+            async 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()()
@@ -174,6 +180,8 @@ async def validate(
         self,
         hostname: str | None,
         connection: ResponseHandler,
+        *,
+        session_manager: SessionManager,
         no_exception: bool = False,
     ) -> (
         list[
@@ -218,20 +226,31 @@ async def validate(
             return None
 
         return await self._validate(
-            hostname, cert_data, telemetry_data, do_retry, no_exception
+            hostname,
+            cert_data,
+            telemetry_data,
+            session_manager=session_manager,
+            do_retry=do_retry,
+            no_exception=no_exception,
         )
 
     async def _validate(
         self,
         hostname: str | None,
         cert_data: list[tuple[Certificate, Certificate]],
         telemetry_data: OCSPTelemetryData,
+        *,
+        session_manager: SessionManager,
         do_retry: bool = True,
         no_exception: bool = False,
     ) -> list[tuple[Exception | None, Certificate, Certificate, CertId, bytes]]:
         """Validate certs sequentially if OCSP response cache server is used."""
         results = await self._validate_certificates_sequential(
-            cert_data, telemetry_data, hostname, do_retry=do_retry
+            cert_data,
+            telemetry_data,
+            hostname=hostname,
+            do_retry=do_retry,
+            session_manager=session_manager,
         )
 
         SnowflakeOCSP.OCSP_CACHE.update_file(self)
@@ -253,6 +272,8 @@ async def _validate_issue_subject(
         issuer: Certificate,
         subject: Certificate,
         telemetry_data: OCSPTelemetryData,
+        *,
+        session_manager: SessionManager,
         hostname: str | None = None,
         do_retry: bool = True,
     ) -> tuple[
@@ -275,7 +296,8 @@ async def _validate_issue_subject(
                 issuer,
                 subject,
                 telemetry_data,
-                hostname,
+                hostname=hostname,
+                session_manager=session_manager,
                 do_retry=do_retry,
                 cache_key=cache_key,
             )
@@ -292,6 +314,8 @@ async def _validate_issue_subject(
     async def _check_ocsp_response_cache_server(
         self,
         cert_data: list[tuple[Certificate, Certificate]],
+        *,
+        session_manager: SessionManager,
     ) -> None:
         """Checks if OCSP response is in cache, and if not it downloads the OCSP response cache from the server.
 
@@ -308,17 +332,23 @@ async def _check_ocsp_response_cache_server(
                 break
 
         if not in_cache:
-            await self.OCSP_CACHE_SERVER.download_cache_from_server(self)
+            await self.OCSP_CACHE_SERVER.download_cache_from_server(
+                self, session_manager=session_manager
+            )
 
     async def _validate_certificates_sequential(
         self,
         cert_data: list[tuple[Certificate, Certificate]],
         telemetry_data: OCSPTelemetryData,
+        *,
+        session_manager: SessionManager,
         hostname: str | None = None,
         do_retry: bool = True,
     ) -> list[tuple[Exception | None, Certificate, Certificate, CertId, bytes]]:
         try:
-            await self._check_ocsp_response_cache_server(cert_data)
+            await self._check_ocsp_response_cache_server(
+                cert_data, session_manager=session_manager
+            )
         except RevocationCheckError as rce:
             telemetry_data.set_event_sub_type(
                 OCSPTelemetryData.ERROR_CODE_MAP[rce.errno]
@@ -339,6 +369,7 @@ async def _validate_certificates_sequential(
                     hostname=hostname,
                     telemetry_data=telemetry_data,
                     do_retry=do_retry,
+                    session_manager=session_manager,
                 )
                 for issuer, subject in cert_data
             ]
@@ -363,6 +394,8 @@ async def validate_by_direct_connection(
         issuer: Certificate,
         subject: Certificate,
         telemetry_data: OCSPTelemetryData,
+        *,
+        session_manager: SessionManager,
         hostname: str = None,
         do_retry: bool = True,
         **kwargs: Any,
@@ -377,7 +410,13 @@ async def validate_by_direct_connection(
                 telemetry_data.set_cache_hit(False)
                 logger.debug("getting OCSP response from CA's OCSP server")
                 ocsp_response = await self._fetch_ocsp_response(
-                    req, subject, cert_id, telemetry_data, hostname, do_retry
+                    req,
+                    subject,
+                    cert_id,
+                    telemetry_data,
+                    session_manager=session_manager,
+                    hostname=hostname,
+                    do_retry=do_retry,
                 )
             else:
                 ocsp_url = self.extract_ocsp_url(subject)
@@ -428,6 +467,8 @@ async def _fetch_ocsp_response(
         subject,
         cert_id,
         telemetry_data,
+        *,
+        session_manager: SessionManager,
         hostname=None,
         do_retry: bool = True,
     ):
@@ -497,7 +538,7 @@ async def _fetch_ocsp_response(
         if not self.is_enabled_fail_open():
             sf_max_retry = SnowflakeOCSP.CA_OCSP_RESPONDER_MAX_RETRY_FC
 
-        async with aiohttp.ClientSession() as session:
+        async 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/aio/_result_batch.py

@@ -13,6 +13,7 @@
     raise_failed_request_error,
     raise_okta_unauthorized_error,
 )
+from snowflake.connector.aio._session_manager import SessionManager
 from snowflake.connector.aio._time_util import TimerContextManager
 from snowflake.connector.arrow_context import ArrowConverterContext
 from snowflake.connector.backoff_policies import exponential_backoff
@@ -111,6 +112,7 @@ def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
                     column_converters,
                     cursor._use_dict_result,
                     json_result_force_utf8_decoding=cursor._connection._json_result_force_utf8_decoding,
+                    session_manager=cursor._connection._session_manager.clone(),
                 )
                 for c in chunks
             ]
@@ -125,6 +127,7 @@ def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
                     cursor._connection._numpy,
                     schema,
                     cursor._connection._arrow_number_to_decimal,
+                    session_manager=cursor._connection._session_manager.clone(),
                 )
                 for c in chunks
             ]
@@ -137,6 +140,7 @@ def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
             schema,
             column_converters,
             cursor._use_dict_result,
+            session_manager=cursor._connection._session_manager.clone(),
         )
     elif rowset_b64 is not None:
         first_chunk = ArrowResultBatch.from_data(
@@ -147,6 +151,7 @@ def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
             cursor._connection._numpy,
             schema,
             cursor._connection._arrow_number_to_decimal,
+            session_manager=cursor._connection._session_manager.clone(),
         )
     else:
         logger.error(f"Don't know how to construct ResultBatches from response: {data}")
@@ -158,6 +163,7 @@ def remote_chunk_info(c: dict[str, Any]) -> RemoteChunkInfo:
             cursor._connection._numpy,
             schema,
             cursor._connection._arrow_number_to_decimal,
+            session_manager=cursor._connection._session_manager.clone(),
         )
 
     return [first_chunk] + rest_of_chunks
@@ -204,7 +210,7 @@ async def _download(
         async def download_chunk(http_session):
             response, content, encoding = None, None, None
             logger.debug(
-                f"downloading result batch id: {self.id} with existing session {http_session}"
+                f"downloading result batch id: {self.id} with session {http_session}"
             )
             response = await http_session.get(**request_data)
             if response.status == OK:
@@ -234,18 +240,29 @@ async def download_chunk(http_session):
                         request_data["timeout"] = aiohttp.ClientTimeout(
                             total=DOWNLOAD_TIMEOUT
                         )
-                    # Try to reuse a connection if possible
-                    if connection and connection._rest is not None:
-                        async with connection._rest._use_session() as session:
+                    # Use SessionManager with same fallback pattern as sync version
+                    if (
+                        connection
+                        and connection.rest
+                        and connection.rest.session_manager is not None
+                    ):
+                        # If connection was explicitly passed and not closed yet - we can reuse SessionManager with session pooling
+                        async with connection.rest.use_session() as session:
                             logger.debug(
                                 f"downloading result batch id: {self.id} with existing session {session}"
                             )
                             response, content, encoding = await download_chunk(session)
+                    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)
+                        async with self._session_manager.use_session() as session:
+                            response, content, encoding = await download_chunk(session)
                     else:
-                        async with aiohttp.ClientSession() as session:
-                            logger.debug(
-                                f"downloading result batch id: {self.id} with new session"
-                            )
+                        # 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
+                        logger.debug(
+                            f"downloading result batch id: {self.id} with new session through local session manager"
+                        )
+                        local_session_manager = SessionManager(use_pooling=False)
+                        async with local_session_manager.use_session() as session:
                             response, content, encoding = await download_chunk(session)
 
                     if response.status == OK: