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

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_session_manager.py

@@ -21,6 +21,7 @@
     resolve_cafile,
 )
 from ._crl import CRLValidator
+from ..url_util import should_bypass_proxies
 from ._ocsp_asn1crypto import SnowflakeOCSPAsn1Crypto
 
 if TYPE_CHECKING:
@@ -439,7 +440,7 @@ def from_config(cls, cfg: AioHttpConfig, **overrides: Any) -> SessionManager:
             cfg = cfg.copy_with(**overrides)
         return cls(config=cfg)
 
-    def make_session(self) -> aiohttp.ClientSession:
+    def make_session(self, *, url: str | None = None) -> aiohttp.ClientSession:
         """Create a new aiohttp.ClientSession with configured connector."""
         connector = self._cfg.get_connector(
             session_manager=self.clone(),
@@ -456,24 +457,18 @@ async def use_session(
         self, url: str | bytes, use_pooling: bool | None = None
     ) -> AsyncGenerator[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()
+            session = self.make_session(url=url)
             try:
                 yield session
             finally:
                 await session.close()
         else:
-            hostname = urlparse(url).hostname if url else None
-            pool = self._sessions_map[hostname]
-            session = pool.get_session()
-            try:
-                yield session
-            finally:
-                pool.return_session(session)
+            with self._yield_session_from_pool(url) as session_from_pool:
+                yield session_from_pool
 
     async def request(
         self,
@@ -585,16 +580,22 @@ def request(
                     )
                 return super().request(method, url, **kwargs)
 
-    def make_session(self) -> aiohttp.ClientSession:
+    def make_session(self, *, url: str | None = None) -> aiohttp.ClientSession:
         connector = self._cfg.get_connector(
             session_manager=self.clone(),
             snowflake_ocsp_mode=self._cfg.snowflake_ocsp_mode,
         )
+        # We use requests.utils here (in asynch code) to keep the behaviour uniform for synch and asynch code. If we wanted each version to depict its http library's behaviour, we could use here: aiohttp.helpers.proxy_bypass(url, proxies={...}) here
+        proxy = (
+            None
+            if should_bypass_proxies(url, no_proxy=self.config.no_proxy)
+            else self.proxy_url
+        )
         # Construct session with base proxy set, request() may override per-URL when bypassing
         return self.SessionWithProxy(
             connector=connector,
             trust_env=self._cfg.trust_env,
-            proxy=self.proxy_url,
+            proxy=proxy,
         )
 
 

src/snowflake/connector/session_manager.py

@@ -11,6 +11,7 @@
 
 from .compat import urlparse
 from .proxy import get_proxy_url
+from .url_util import should_bypass_proxies
 from .vendored import requests
 from .vendored.requests import Response, Session
 from .vendored.requests.adapters import BaseAdapter, HTTPAdapter
@@ -610,7 +611,7 @@ def make_session(self, *, url: str | None = None) -> Session:
             {
                 "no_proxy": self._cfg.no_proxy,
             }
-            if requests.utils.should_bypass_proxies(url, no_proxy=self.config.no_proxy)
+            if should_bypass_proxies(url, no_proxy=self.config.no_proxy)
             else {
                 "http": self.proxy_url,
                 "https": self.proxy_url,

src/snowflake/connector/url_util.py

@@ -3,6 +3,9 @@
 import re
 import urllib.parse
 from logging import getLogger
+from typing import Iterable
+
+import vendored.requests as requests
 
 from .constants import _TOP_LEVEL_DOMAIN_REGEX
 
@@ -47,3 +50,7 @@ def extract_top_level_domain_from_hostname(hostname: str | None = None) -> str:
     # RFC1034 for TLD spec, and https://data.iana.org/TLD/tlds-alpha-by-domain.txt for full TLD list
     match = re.search(_TOP_LEVEL_DOMAIN_REGEX, hostname)
     return (match.group(0)[1:] if match else "com").lower()
+
+
+def should_bypass_proxies(url: str | bytes, no_proxy: Iterable[str] | None) -> bool:
+    return requests.utils.should_bypass_proxies(url, no_proxy)

test/unit/aio/mock_utils.py

@@ -34,8 +34,8 @@ async def forbidden_connect(*args, **kwargs):
         raise NotImplementedError("Unit test tried to make real network connection")
 
     class MockSessionManager(SessionManager):
-        def make_session(self):
-            session = super().make_session()
+        def make_session(self, *, url: str | None = None):
+            session = super().make_session(url)
             if not allow_send:
                 # Block at connector._connect level (like sync blocks session.send)
                 # This allows patches on session.request to work