[async] Applied #2596 to async code - part 1

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_connection.py

@@ -1047,6 +1047,15 @@ async def connect(self, **kwargs) -> None:
             proxy_password=self.proxy_password,
             snowflake_ocsp_mode=self._ocsp_mode(),
             trust_env=True,  # Required for proxy support via environment variables
+            no_proxy=(
+                ",".join(str(x) for x in self.no_proxy)
+                if (
+                    self.no_proxy is not None
+                    and isinstance(self.no_proxy, Iterable)
+                    and not isinstance(self.no_proxy, (str, bytes))
+                )
+                else self.no_proxy
+            ),
         )
         self._session_manager = SessionManagerFactory.get_manager(self._http_config)
 

src/snowflake/connector/aio/_ocsp_snowflake.py

@@ -98,7 +98,7 @@ async def _download_ocsp_response_cache(
                 if sf_cache_server_url is not None:
                     url = sf_cache_server_url
 
-            async with session_manager.use_session() as session:
+            async with session_manager.use_session(url) as session:
                 max_retry = SnowflakeOCSP.OCSP_CACHE_SERVER_MAX_RETRY if do_retry else 1
                 sleep_time = 1
                 backoff = exponential_backoff()()
@@ -544,7 +544,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 session_manager.use_session() as session:
+        async with session_manager.use_session(target_url) 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

@@ -247,14 +247,18 @@ async def download_chunk(http_session):
                         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:
+                        async with connection.rest.use_session(
+                            request_data["url"]
+                        ) 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:
+                        async with self._session_manager.use_session(
+                            request_data["url"]
+                        ) as session:
                             response, content, encoding = await download_chunk(session)
                     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
@@ -264,7 +268,9 @@ async def download_chunk(http_session):
                         local_session_manager = SessionManagerFactory.get_manager(
                             use_pooling=False
                         )
-                        async with local_session_manager.use_session() as session:
+                        async with local_session_manager.use_session(
+                            request_data["url"]
+                        ) as session:
                             response, content, encoding = await download_chunk(session)
 
                     if response.status == OK:

src/snowflake/connector/aio/_session_manager.py

@@ -257,7 +257,7 @@ class _RequestVerbsUsingSessionMixin(abc.ABC):
 
     @abc.abstractmethod
     async def use_session(
-        self, url: str, use_pooling: bool
+        self, url: str | bytes, use_pooling: bool
     ) -> AsyncGenerator[aiohttp.ClientSession]: ...
 
     async def get(
@@ -453,9 +453,12 @@ def make_session(self) -> aiohttp.ClientSession:
 
     @contextlib.asynccontextmanager
     async def use_session(
-        self, url: str | bytes | None = None, use_pooling: bool | None = None
+        self, url: str | bytes, use_pooling: bool | None = None
     ) -> AsyncGenerator[aiohttp.ClientSession]:
-        """Async version of use_session yielding aiohttp.ClientSession."""
+        """
+        Async version of use_session yielding aiohttp.ClientSession.
+        'url' is an obligatory parameter due to the need for correct proxy handling (i.e. bypassing caused by no_proxy settings).
+        """
         use_pooling = use_pooling if use_pooling is not None else self.use_pooling
         if not use_pooling:
             session = self.make_session()

src/snowflake/connector/aio/_wif_util.py

@@ -14,6 +14,7 @@
 
 from ..errorcode import ER_WIF_CREDENTIALS_NOT_FOUND
 from ..errors import MissingDependencyError, ProgrammingError
+from ..session_manager import SessionManagerFactory
 from ..wif_util import (
     DEFAULT_ENTRA_SNOWFLAKE_RESOURCE,
     SNOWFLAKE_AUDIENCE,
@@ -307,7 +308,7 @@ async def create_attestation(
     session_manager = (
         session_manager.clone()
         if session_manager
-        else SessionManager(use_pooling=True, max_retries=0)
+        else SessionManagerFactory.get_manager(use_pooling=True, max_retries=0)
     )
 
     if provider == AttestationProvider.AWS: