[async] Applied #2451 to async code - part 3 - test passing

Author: sfc-gh-fpawlowski

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 SessionManagerFactory
 from snowflake.connector.aio._time_util import TimerContextManager
 from snowflake.connector.arrow_context import ArrowConverterContext
 from snowflake.connector.backoff_policies import exponential_backoff
@@ -33,7 +34,6 @@
 from snowflake.connector.result_batch import ResultBatch as ResultBatchSync
 from snowflake.connector.result_batch import _create_nanoarrow_iterator
 from snowflake.connector.secret_detector import SecretDetector
-from src.snowflake.connector.aio._session_manager import SessionManagerFactory
 
 if TYPE_CHECKING:
     from pandas import DataFrame

src/snowflake/connector/aio/_session_manager.py

@@ -441,7 +441,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(
@@ -480,7 +480,7 @@ async def request(
             url: StrOrURL,
             **kwargs: Unpack[_RequestOptions],
         ):
-            # Apply proxy per request and inject Host header when proxying
+            # Inject Host header when proxying
             try:
                 # respect caller-provided proxy and proxy_headers if any
                 provided_proxy = kwargs.get("proxy") or self._default_proxy
@@ -515,21 +515,6 @@ def make_session(self) -> aiohttp.ClientSession:
             proxy=self.proxy_url,
         )
 
-    def clone(
-        self,
-        *,
-        use_pooling: bool | None = None,
-        connector_factory: ConnectorFactory | None = None,
-    ) -> SessionManager:
-        """Return a new async SessionManager sharing this instance's config."""
-        overrides: dict[str, Any] = {}
-        if use_pooling is not None:
-            overrides["use_pooling"] = use_pooling
-        if connector_factory is not None:
-            overrides["connector_factory"] = connector_factory
-
-        return SessionManager.from_config(self._cfg, **overrides)
-
 
 class SessionManagerFactory:
     @staticmethod

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": [

test/unit/aio/test_connection_async_unit.py

@@ -27,6 +27,7 @@
 from cryptography.hazmat.primitives.asymmetric import rsa
 
 import snowflake.connector.aio
+from snowflake.connector.aio import connect as async_connect
 from snowflake.connector.aio._network import SnowflakeRestful
 from snowflake.connector.aio.auth import (
     AuthByDefault,
@@ -773,3 +774,93 @@ async def test_invalid_authenticator():
         )
         await conn.connect()
     assert "Unknown authenticator: INVALID" in str(excinfo.value)
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize("proxy_method", ["explicit_args", "env_vars"])
+async def test_large_query_through_proxy_async(
+    wiremock_generic_mappings_dir,
+    wiremock_target_proxy_pair,
+    wiremock_mapping_dir,
+    proxy_env_vars,
+    proxy_method,
+):
+    target_wm, proxy_wm = wiremock_target_proxy_pair
+
+    password_mapping = wiremock_mapping_dir / "auth/password/successful_flow.json"
+    multi_chunk_request_mapping = (
+        wiremock_mapping_dir / "queries/select_large_request_successful.json"
+    )
+    disconnect_mapping = (
+        wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json"
+    )
+    telemetry_mapping = wiremock_generic_mappings_dir / "telemetry.json"
+    chunk_1_mapping = wiremock_mapping_dir / "queries/chunk_1.json"
+    chunk_2_mapping = wiremock_mapping_dir / "queries/chunk_2.json"
+
+    expected_headers = {"Via": {"contains": "wiremock"}}
+
+    target_wm.import_mapping(password_mapping, expected_headers=expected_headers)
+    target_wm.add_mapping_with_default_placeholders(
+        multi_chunk_request_mapping, expected_headers
+    )
+    target_wm.add_mapping(disconnect_mapping, expected_headers=expected_headers)
+    target_wm.add_mapping(telemetry_mapping, expected_headers=expected_headers)
+    target_wm.add_mapping_with_default_placeholders(chunk_1_mapping, expected_headers)
+    target_wm.add_mapping_with_default_placeholders(chunk_2_mapping, expected_headers)
+
+    set_proxy_env_vars, clear_proxy_env_vars = proxy_env_vars
+    connect_kwargs = {
+        "user": "testUser",
+        "password": "testPassword",
+        "account": "testAccount",
+        "host": target_wm.wiremock_host,
+        "port": target_wm.wiremock_http_port,
+        "protocol": "http",
+        "warehouse": "TEST_WH",
+    }
+
+    if proxy_method == "explicit_args":
+        connect_kwargs.update(
+            {
+                "proxy_host": proxy_wm.wiremock_host,
+                "proxy_port": str(proxy_wm.wiremock_http_port),
+                "proxy_user": "proxyUser",
+                "proxy_password": "proxyPass",
+            }
+        )
+        clear_proxy_env_vars()
+    else:
+        proxy_url = f"http://proxyUser:proxyPass@{proxy_wm.wiremock_host}:{proxy_wm.wiremock_http_port}"
+        set_proxy_env_vars(proxy_url)
+
+    row_count = 50_000
+    conn = await async_connect(**connect_kwargs)
+    try:
+        cur = conn.cursor()
+        await cur.execute(
+            f"select seq4() as n from table(generator(rowcount => {row_count}));"
+        )
+        assert len(cur._result_set.batches) > 1
+        _ = [r async for r in cur]
+    finally:
+        await conn.close()
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(
+            f"{proxy_wm.http_host_with_port}/__admin/requests"
+        ) as resp:
+            proxy_reqs = await resp.json()
+        assert any(
+            "/queries/v1/query-request" in r["request"]["url"]
+            for r in proxy_reqs["requests"]
+        )
+
+        async with session.get(
+            f"{target_wm.http_host_with_port}/__admin/requests"
+        ) as resp:
+            target_reqs = await resp.json()
+        assert any(
+            "/queries/v1/query-request" in r["request"]["url"]
+            for r in target_reqs["requests"]
+        )

test/unit/aio/test_oauth_token_async.py

@@ -7,6 +7,7 @@
 from unittest import mock
 from unittest.mock import Mock, patch
 
+import aiohttp
 import pytest
 
 try:
@@ -18,7 +19,7 @@
 import snowflake.connector.errors
 from snowflake.connector.token_cache import TokenCache, TokenKey, TokenType
 
-from ...wiremock.wiremock_utils import WiremockClient
+from ...test_utils.wiremock.wiremock_utils import WiremockClient
 from ..test_oauth_token import omit_oauth_urls_check  # noqa: F401
 
 logger = logging.getLogger(__name__)
@@ -699,3 +700,159 @@ async def test_client_creds_expired_refresh_token_flow_async(
     cached_refresh_token = temp_cache_async.retrieve(refresh_token_key)
     assert cached_access_token == "expired-access-token-123"
     assert cached_refresh_token == "expired-refresh-token-123"
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize("proxy_method", ["explicit_args", "env_vars"])
+async def test_client_credentials_flow_through_proxy_async(
+    wiremock_oauth_client_creds_dir,
+    wiremock_generic_mappings_dir,
+    wiremock_target_proxy_pair,
+    temp_cache_async,
+    proxy_env_vars,
+    proxy_method,
+):
+    """Run OAuth Client-Credentials flow and ensure it goes through proxy (async)."""
+    from snowflake.connector.aio import SnowflakeConnection
+
+    target_wm, proxy_wm = wiremock_target_proxy_pair
+
+    expected_headers = {"Via": {"contains": "wiremock"}}
+
+    target_wm.import_mapping_with_default_placeholders(
+        wiremock_oauth_client_creds_dir / "successful_flow.json", expected_headers
+    )
+    target_wm.add_mapping_with_default_placeholders(
+        wiremock_generic_mappings_dir / "snowflake_login_successful.json",
+        expected_headers,
+    )
+    target_wm.add_mapping(
+        wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json",
+        expected_headers=expected_headers,
+    )
+
+    token_request_url = f"http://{target_wm.wiremock_host}:{target_wm.wiremock_http_port}/oauth/token-request"
+
+    set_proxy_env_vars, clear_proxy_env_vars = proxy_env_vars
+    connect_kwargs = {
+        "user": "testUser",
+        "authenticator": "OAUTH_CLIENT_CREDENTIALS",
+        "oauth_client_id": "cid",
+        "oauth_client_secret": "secret",
+        "account": "testAccount",
+        "protocol": "http",
+        "role": "ANALYST",
+        "oauth_token_request_url": token_request_url,
+        "host": target_wm.wiremock_host,
+        "port": target_wm.wiremock_http_port,
+        "oauth_enable_refresh_tokens": True,
+        "client_store_temporary_credential": True,
+        "token_cache": temp_cache_async,
+    }
+
+    if proxy_method == "explicit_args":
+        connect_kwargs.update(
+            {
+                "proxy_host": proxy_wm.wiremock_host,
+                "proxy_port": str(proxy_wm.wiremock_http_port),
+                "proxy_user": "proxyUser",
+                "proxy_password": "proxyPass",
+            }
+        )
+        clear_proxy_env_vars()
+    else:
+        proxy_url = f"http://proxyUser:proxyPass@{proxy_wm.wiremock_host}:{proxy_wm.wiremock_http_port}"
+        set_proxy_env_vars(proxy_url)
+
+    with mock.patch("secrets.token_urlsafe", return_value="abc123"):
+        cnx = SnowflakeConnection(**connect_kwargs)
+        await cnx.connect()
+        await cnx.close()
+
+    async with aiohttp.ClientSession() as session:
+        async with session.get(
+            f"{proxy_wm.http_host_with_port}/__admin/requests"
+        ) as resp:
+            proxy_requests = await resp.json()
+        assert any(
+            req["request"]["url"].endswith("/oauth/token-request")
+            for req in proxy_requests["requests"]
+        )
+
+        async with session.get(
+            f"{target_wm.http_host_with_port}/__admin/requests"
+        ) as resp:
+            target_requests = await resp.json()
+        assert any(
+            req["request"]["url"].endswith("/oauth/token-request")
+            for req in target_requests["requests"]
+        )
+
+
+@pytest.mark.skipolddriver
+@patch("snowflake.connector.auth._http_server.AuthHttpServer.DEFAULT_TIMEOUT", 30)
+async def test_oauth_code_successful_flow_through_proxy_async(
+    wiremock_oauth_authorization_code_dir,
+    wiremock_generic_mappings_dir,
+    wiremock_target_proxy_pair,
+    webbrowser_mock_sync,
+    monkeypatch,
+    omit_oauth_urls_check,  # noqa: F811
+) -> None:
+    from snowflake.connector.aio import SnowflakeConnection
+
+    monkeypatch.setenv("SNOWFLAKE_AUTH_SOCKET_REUSE_PORT", "true")
+    target_wm, proxy_wm = wiremock_target_proxy_pair
+
+    target_wm.import_mapping_with_default_placeholders(
+        wiremock_oauth_authorization_code_dir / "successful_flow.json",
+    )
+    target_wm.add_mapping_with_default_placeholders(
+        wiremock_generic_mappings_dir / "snowflake_login_successful.json",
+    )
+    target_wm.add_mapping(
+        wiremock_generic_mappings_dir / "snowflake_disconnect_successful.json",
+    )
+
+    with mock.patch("webbrowser.open", new=webbrowser_mock_sync.open):
+        with mock.patch("secrets.token_urlsafe", return_value="abc123"):
+            cnx = SnowflakeConnection(
+                user="testUser",
+                authenticator="OAUTH_AUTHORIZATION_CODE",
+                oauth_client_id="123",
+                account="testAccount",
+                protocol="http",
+                role="ANALYST",
+                proxy_host=proxy_wm.wiremock_host,
+                proxy_port=str(proxy_wm.wiremock_http_port),
+                proxy_user="proxyUser",
+                proxy_password="proxyPass",
+                oauth_client_secret="testClientSecret",
+                oauth_token_request_url=f"http://{target_wm.wiremock_host}:{target_wm.wiremock_http_port}/oauth/token-request",
+                oauth_authorization_url=f"http://{target_wm.wiremock_host}:{target_wm.wiremock_http_port}/oauth/authorize",
+                oauth_redirect_uri="http://localhost:8009/snowflake/oauth-redirect",
+                host=target_wm.wiremock_host,
+                port=target_wm.wiremock_http_port,
+            )
+
+            await cnx.connect()
+            await cnx.close()
+
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                f"http://{proxy_wm.wiremock_host}:{proxy_wm.wiremock_http_port}/__admin/requests"
+            ) as resp:
+                proxy_requests = await resp.json()
+            assert any(
+                req["request"]["url"].endswith("/oauth/token-request")
+                for req in proxy_requests["requests"]
+            ), "Proxy did not record token-request"
+
+            async with session.get(
+                f"http://{target_wm.wiremock_host}:{target_wm.wiremock_http_port}/__admin/requests"
+            ) as resp:
+                target_requests = await resp.json()
+            assert any(
+                req["request"]["url"].endswith("/oauth/token-request")
+                for req in target_requests["requests"]
+            ), "Target did not receive token-request forwarded by proxy"

test/unit/aio/test_programmatic_access_token_async.py

@@ -16,7 +16,7 @@
 
 import snowflake.connector.errors
 
-from ...wiremock.wiremock_utils import WiremockClient
+from ...test_utils.wiremock import WiremockClient
 
 
 @pytest.mark.skipolddriver