[async] Applied #2451 to async code - test passing, ProxySessionManager for async with SessionWithProxy

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_connection.py

@@ -21,7 +21,6 @@
     Error,
     OperationalError,
     ProgrammingError,
-    proxy,
 )
 
 from .._query_context_cache import QueryContextCache
@@ -80,6 +79,7 @@
 from ._session_manager import (
     AioHttpConfig,
     SessionManager,
+    SessionManagerFactory,
     SnowflakeSSLConnectorFactory,
 )
 from ._telemetry import TelemetryClient
@@ -191,10 +191,6 @@ async def __open_connection(self):
             use_numpy=self._numpy, support_negative_year=self._support_negative_year
         )
 
-        proxy.set_proxies(
-            self.proxy_host, self.proxy_port, self.proxy_user, self.proxy_password
-        )
-
         self._rest = SnowflakeRestful(
             host=self.host,
             port=self.port,
@@ -1014,13 +1010,17 @@ async def connect(self, **kwargs) -> None:
         else:
             self.__config(**self._conn_parameters)
 
-        self._http_config = AioHttpConfig(
+        self._http_config: AioHttpConfig = AioHttpConfig(
             connector_factory=SnowflakeSSLConnectorFactory(),
             use_pooling=not self.disable_request_pooling,
+            proxy_host=self.proxy_host,
+            proxy_port=self.proxy_port,
+            proxy_user=self.proxy_user,
+            proxy_password=self.proxy_password,
             snowflake_ocsp_mode=self._ocsp_mode(),
             trust_env=True,  # Required for proxy support via environment variables
         )
-        self._session_manager = SessionManager(self._http_config)
+        self._session_manager = SessionManagerFactory.get_manager(self._http_config)
 
         if self.enable_connection_diag:
             raise NotImplementedError(

src/snowflake/connector/aio/_network.py

@@ -10,7 +10,6 @@
 from typing import TYPE_CHECKING, Any, AsyncGenerator
 
 import OpenSSL.SSL
-from urllib3.util.url import parse_url
 
 from ..compat import FORBIDDEN, OK, UNAUTHORIZED, urlencode, urlparse, urlsplit
 from ..constants import (
@@ -79,7 +78,11 @@
 )
 from ..time_util import TimeoutBackoffCtx
 from ._description import CLIENT_NAME
-from ._session_manager import SessionManager, SnowflakeSSLConnectorFactory
+from ._session_manager import (
+    SessionManager,
+    SessionManagerFactory,
+    SnowflakeSSLConnectorFactory,
+)
 
 if TYPE_CHECKING:
     from snowflake.connector.aio import SnowflakeConnection
@@ -145,15 +148,12 @@ def __init__(
             session_manager = (
                 connection._session_manager
                 if (connection and connection._session_manager)
-                else SessionManager(connector_factory=SnowflakeSSLConnectorFactory())
+                else SessionManagerFactory.get_manager(
+                    connector_factory=SnowflakeSSLConnectorFactory()
+                )
             )
         self._session_manager = session_manager
 
-        if self._connection and self._connection.proxy_host:
-            self._get_proxy_headers = lambda url: {"Host": parse_url(url).hostname}
-        else:
-            self._get_proxy_headers = lambda _: None
-
     async def close(self) -> None:
         if hasattr(self, "_token"):
             del self._token
@@ -737,7 +737,6 @@ async def _request_exec(
                 headers=headers,
                 data=input_data,
                 timeout=aiohttp.ClientTimeout(socket_timeout),
-                proxy_headers=self._get_proxy_headers(full_url),
             )
             try:
                 if raw_ret.status == OK:

src/snowflake/connector/aio/_result_batch.py

@@ -13,7 +13,7 @@
     raise_failed_request_error,
     raise_okta_unauthorized_error,
 )
-from snowflake.connector.aio._session_manager import SessionManager
+from snowflake.connector.aio._session_manager import SessionManagerFactory
 from snowflake.connector.aio._time_util import TimerContextManager
 from snowflake.connector.arrow_context import ArrowConverterContext
 from snowflake.connector.backoff_policies import exponential_backoff
@@ -261,7 +261,9 @@ async def download_chunk(http_session):
                         logger.debug(
                             f"downloading result batch id: {self.id} with new session through local session manager"
                         )
-                        local_session_manager = SessionManager(use_pooling=False)
+                        local_session_manager = SessionManagerFactory.get_manager(
+                            use_pooling=False
+                        )
                         async with local_session_manager.use_session() as session:
                             response, content, encoding = await download_chunk(session)
 

src/snowflake/connector/aio/_session_manager.py

@@ -4,8 +4,10 @@
 from typing import TYPE_CHECKING
 
 from aiohttp import ClientRequest, ClientTimeout
+from aiohttp.client import _RequestOptions
 from aiohttp.client_proto import ResponseHandler
 from aiohttp.connector import Connection
+from aiohttp.typedefs import StrOrURL
 
 from .. import OperationalError
 from ..errorcode import ER_OCSP_RESPONSE_CERT_STATUS_REVOKED
@@ -14,6 +16,8 @@
 
 if TYPE_CHECKING:
     from aiohttp.tracing import Trace
+    from typing import Unpack
+    from aiohttp.client import _RequestContextManager
 
 import abc
 import collections
@@ -44,10 +48,10 @@ def __init__(
     ):
         self._snowflake_ocsp_mode = snowflake_ocsp_mode
         if session_manager is None:
-            logger.debug(
-                "SessionManager instance was not passed to SSLConnector - OCSP will use default settings which may be distinct from the customer's specific one. Code should always pass such instance so please verify why it isn't true in the current context"
+            logger.warning(
+                "SessionManager instance was not passed to SSLConnector - OCSP will use default settings which may be distinct from the customer's specific one. Code should always pass such instance - verify why it isn't true in the current context"
             )
-            session_manager = SessionManager()
+            session_manager = SessionManagerFactory.get_manager()
         self._session_manager = session_manager
         if self._snowflake_ocsp_mode == OCSPMode.FAIL_OPEN and sys.version_info < (
             3,
@@ -345,24 +349,38 @@ def __init__(
             lambda: SessionPool(self)
         )
 
+    @classmethod
+    def from_config(cls, cfg: AioHttpConfig, **overrides: Any) -> SessionManager:
+        """Build a new manager from *cfg*, optionally overriding fields.
+
+        Example::
+
+            no_pool_cfg = conn._http_config.copy_with(use_pooling=False)
+            manager = SessionManager.from_config(no_pool_cfg)
+        """
+
+        if overrides:
+            cfg = cfg.copy_with(**overrides)
+        return cls(config=cfg)
+
     @property
     def connector_factory(self) -> Callable[..., aiohttp.BaseConnector]:
         return self._cfg.connector_factory
 
     @connector_factory.setter
     def connector_factory(self, value: Callable[..., aiohttp.BaseConnector]) -> None:
-        self._cfg = self._cfg.copy_with(connector_factory=value)
+        self._cfg: AioHttpConfig = self._cfg.copy_with(connector_factory=value)
 
     def make_session(self) -> aiohttp.ClientSession:
         """Create a new aiohttp.ClientSession with configured connector."""
         connector = self._cfg.get_connector(
             session_manager=self.clone(),
             snowflake_ocsp_mode=self._cfg.snowflake_ocsp_mode,
         )
-
         return aiohttp.ClientSession(
             connector=connector,
             trust_env=self._cfg.trust_env,
+            proxy=self.proxy_url,
         )
 
     @contextlib.asynccontextmanager
@@ -425,7 +443,7 @@ def clone(
         if connector_factory is not None:
             overrides["connector_factory"] = connector_factory
 
-        return SessionManager.from_config(self._cfg, **overrides)
+        return self.from_config(self._cfg, **overrides)
 
 
 async def request(
@@ -454,3 +472,82 @@ async def request(
         use_pooling=use_pooling,
         **kwargs,
     )
+
+
+class ProxySessionManager(SessionManager):
+    class SessionWithProxy(aiohttp.ClientSession):
+        if sys.version_info >= (3, 11) and TYPE_CHECKING:
+
+            def request(
+                self,
+                method: str,
+                url: StrOrURL,
+                **kwargs: Unpack[_RequestOptions],
+            ) -> _RequestContextManager: ...
+
+        else:
+
+            def request(
+                self, method: str, url: StrOrURL, **kwargs: Any
+            ) -> _RequestContextManager:
+                """Perform HTTP request."""
+                # Inject Host header when proxying
+                try:
+                    # respect caller-provided proxy and proxy_headers if any
+                    provided_proxy = kwargs.get("proxy") or self._default_proxy
+                    provided_proxy_headers = kwargs.get("proxy_headers")
+                    if provided_proxy is not None:
+                        authority = urlparse(str(url)).netloc
+                        if provided_proxy_headers is None:
+                            kwargs["proxy_headers"] = {"Host": authority}
+                        elif "Host" not in provided_proxy_headers:
+                            provided_proxy_headers["Host"] = authority
+                        else:
+                            logger.debug(
+                                "Host header was already set - not overriding with netloc at the ClientSession.request method level."
+                            )
+                except Exception:
+                    logger.warning(
+                        "Failed to compute proxy settings for %s",
+                        urlparse(url).hostname,
+                        exc_info=True,
+                    )
+                return super().request(method, url, **kwargs)
+
+    def make_session(self) -> aiohttp.ClientSession:
+        connector = self._cfg.get_connector(
+            session_manager=self.clone(),
+            snowflake_ocsp_mode=self._cfg.snowflake_ocsp_mode,
+        )
+        # 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,
+        )
+
+
+class SessionManagerFactory:
+    @staticmethod
+    def get_manager(
+        config: AioHttpConfig | None = None, **http_config_kwargs
+    ) -> SessionManager:
+        """Return a proxy-aware or plain async SessionManager based on config.
+
+        If any explicit proxy parameters are provided (in config or kwargs),
+        return ProxySessionManager; otherwise return the base SessionManager.
+        """
+
+        def _has_proxy_params(cfg: AioHttpConfig | None, kwargs: dict) -> bool:
+            cfg_keys = (
+                "proxy_host",
+                "proxy_port",
+            )
+            in_cfg = any(getattr(cfg, k, None) for k in cfg_keys) if cfg else False
+            in_kwargs = "proxy" in kwargs
+            return in_cfg or in_kwargs
+
+        if _has_proxy_params(config, http_config_kwargs):
+            return ProxySessionManager(config, **http_config_kwargs)
+        else:
+            return SessionManager(config, **http_config_kwargs)

src/snowflake/connector/aio/_storage_client.py

@@ -15,7 +15,7 @@
 from ..encryption_util import SnowflakeEncryptionUtil
 from ..errors import RequestExceedMaxRetryError
 from ..storage_client import SnowflakeStorageClient as SnowflakeStorageClientSync
-from ._session_manager import SessionManager
+from ._session_manager import SessionManagerFactory
 
 if TYPE_CHECKING:  # pragma: no cover
     from ..file_transfer_agent import SnowflakeFileMeta, StorageCredential
@@ -205,7 +205,9 @@ async def _send_request_with_retry(
                     # SessionManager on the fly, if code ends up here, since we probably do not care about losing
                     # proxy or HTTP setup.
                     logger.debug("storage client request with new session")
-                    session_manager = SessionManager(use_pooling=False)
+                    session_manager = SessionManagerFactory.get_manager(
+                        use_pooling=False
+                    )
                     response = await session_manager.request(verb, url, **rest_kwargs)
 
                 if await self._has_expired_presigned_url(response):

src/snowflake/connector/aio/_wif_util.py

@@ -21,7 +21,7 @@
     extract_iss_and_sub_without_signature_verification,
     get_aws_sts_hostname,
 )
-from ._session_manager import SessionManager
+from ._session_manager import SessionManager, SessionManagerFactory
 
 logger = logging.getLogger(__name__)
 
@@ -187,7 +187,9 @@ async def create_attestation(
     """
     entra_resource = entra_resource or DEFAULT_ENTRA_SNOWFLAKE_RESOURCE
     session_manager = (
-        session_manager.clone() if session_manager else SessionManager(use_pooling=True)
+        session_manager.clone()
+        if session_manager
+        else SessionManagerFactory.get_manager(use_pooling=True)
     )
 
     if provider == AttestationProvider.AWS:

src/snowflake/connector/result_batch.py

@@ -261,7 +261,7 @@ def __init__(
             [s._to_result_metadata_v1() for s in schema] if schema is not None else None
         )
         self._use_dict_result = use_dict_result
-        # Passed to contain the configured Http behavior in case the connectio is no longer active for the download
+        # Passed to contain the configured Http behavior in case the connection is no longer active for the download
         # Can be overridden with setters if needed.
         self._session_manager = session_manager
         self._metrics: dict[str, int] = {}

test/data/wiremock/mappings/auth/password/successful_flow.json

@@ -18,6 +18,7 @@
       },
       "response": {
         "status": 200,
+        "headers": { "Content-Type": "application/json" },
         "jsonBody": {
           "data": {
             "masterToken": "master token",

test/data/wiremock/mappings/queries/select_1_successful.json

@@ -11,6 +11,7 @@
   },
   "response": {
     "status": 200,
+    "headers": { "Content-Type": "application/json" },
     "jsonBody": {
       "data": {
         "parameters": [

test/data/wiremock/mappings/queries/select_large_request_successful.json

@@ -11,6 +11,7 @@
   },
   "response": {
     "status": 200,
+    "headers": { "Content-Type": "application/json" },
     "jsonBody": {
       "data": {
         "parameters": [