[async] Applied #2596 to async code - part 3 - basic tests working

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_session_manager.py

@@ -467,7 +467,7 @@ async def use_session(
             finally:
                 await session.close()
         else:
-            with self._yield_session_from_pool(url) as session_from_pool:
+            for session_from_pool in self._yield_session_from_pool(url):
                 yield session_from_pool
 
     async def request(

src/snowflake/connector/url_util.py

@@ -5,9 +5,8 @@
 from logging import getLogger
 from typing import Iterable
 
-import vendored.requests as requests
-
 from .constants import _TOP_LEVEL_DOMAIN_REGEX
+from .vendored import requests
 
 logger = getLogger(__name__)
 

test/unit/aio/test_connection_async_unit.py

@@ -26,7 +26,6 @@
 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,
@@ -842,93 +841,3 @@ 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_proxies_async.py

@@ -4,6 +4,7 @@
 import pytest
 
 from snowflake.connector.aio import connect
+from snowflake.connector.aio import connect as async_connect
 
 pytestmark = pytest.mark.asyncio
 
@@ -86,3 +87,93 @@ async def test_basic_query_through_proxy_async(
             "/queries/v1/query-request" in r["request"]["url"]
             for r in target_reqs["requests"]
         )
+
+
+@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"]
+        )