[async] Applied #2489 to async code

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_session_manager.py

@@ -34,6 +34,7 @@
 from ..session_manager import BaseHttpConfig
 from ..session_manager import SessionManager as SessionManagerSync
 from ..session_manager import SessionPool as SessionPoolSync
+from ..session_manager import _ConfigDirectAccessMixin
 
 logger = logging.getLogger(__name__)
 
@@ -328,7 +329,29 @@ async def delete(
             )
 
 
-class SessionManager(_RequestVerbsUsingSessionMixin, SessionManagerSync):
+class _AsyncHttpConfigDirectAccessMixin(_ConfigDirectAccessMixin, abc.ABC):
+    @property
+    @abc.abstractmethod
+    def config(self) -> AioHttpConfig: ...
+
+    @config.setter
+    @abc.abstractmethod
+    def config(self, value) -> AioHttpConfig: ...
+
+    @property
+    def connector_factory(self) -> Callable[..., aiohttp.BaseConnector]:
+        return self.config.connector_factory
+
+    @connector_factory.setter
+    def connector_factory(self, value: Callable[..., aiohttp.BaseConnector]) -> None:
+        self.config: AioHttpConfig = self.config.copy_with(connector_factory=value)
+
+
+class SessionManager(
+    _RequestVerbsUsingSessionMixin,
+    SessionManagerSync,
+    _AsyncHttpConfigDirectAccessMixin,
+):
     """
     Async HTTP session manager for aiohttp.ClientSession instances.
 
@@ -363,14 +386,6 @@ def from_config(cls, cfg: AioHttpConfig, **overrides: Any) -> SessionManager:
             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: 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(
@@ -432,18 +447,18 @@ async def close(self):
 
     def clone(
         self,
-        *,
-        use_pooling: bool | None = None,
-        connector_factory: ConnectorFactory | None = None,
+        **http_config_overrides,
     ) -> 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 self.from_config(self._cfg, **overrides)
+        """Return a new *stateless* SessionManager sharing this instance’s config.
+
+        "Shallow clone" - the configuration object (HttpConfig) is reused as-is,
+        while *stateful* aspects such as the per-host SessionPool mapping are
+        reset, so the two managers do not share live `requests.Session`
+        objects.
+        Optional kwargs (e.g. *use_pooling* / *adapter_factory* / max_retries etc.) - overrides to create a modified
+        copy of the HttpConfig before instantiation.
+        """
+        return self.from_config(self._cfg, **http_config_overrides)
 
 
 async def request(

test/integ/aio_it/test_connection_async.py

@@ -58,6 +58,13 @@
 except ImportError:
     pass
 
+from test.integ.test_connection import (
+    _assert_log_bytes_within_tolerance,
+    _calculate_log_bytes,
+    _find_matching_patterns,
+    _log_pattern_analysis,
+)
+
 
 async def test_basic(conn_testaccount):
     """Basic Connection test."""
@@ -1614,3 +1621,85 @@ async def test_snowflake_version():
     assert re.match(
         version_pattern, await conn.snowflake_version
     ), f"snowflake_version should match pattern 'x.y.z', but got '{await conn.snowflake_version}'"
+
+
+@pytest.mark.skipolddriver
+async def test_logs_size_during_basic_query_stays_unchanged(conn_cnx, caplog):
+    """Test that the amount of bytes logged during normal select 1 flow is within acceptable range. Related to: SNOW-2268606"""
+    caplog.set_level(logging.INFO, "snowflake.connector")
+    caplog.clear()
+
+    # Test-specific constants
+    EXPECTED_BYTES = 145
+    ACCEPTABLE_DELTA = 0.6
+    EXPECTED_PATTERNS = [
+        "Snowflake Connector for Python Version: ",  # followed by version info
+        "Connecting to GLOBAL Snowflake domain",
+    ]
+
+    async with conn_cnx() as conn:
+        async with conn.cursor() as cur:
+            await (await cur.execute("select 1")).fetchall()
+
+            actual_messages = [record.getMessage() for record in caplog.records]
+            total_log_bytes = _calculate_log_bytes(actual_messages)
+
+            if total_log_bytes != EXPECTED_BYTES:
+                logging.warning(
+                    f"There was a change in a size of the logs produced by the basic Snowflake query. "
+                    f"Expected: {EXPECTED_BYTES}, got: {total_log_bytes}. "
+                    f"We may need to update the test_logs_size_during_basic_query_stays_unchanged - i.e. EXACT_EXPECTED_LOGS_BYTES constant."
+                )
+
+                # Check if patterns match to decide whether to show all messages
+                matched_patterns, missing_patterns, unmatched_messages = (
+                    _find_matching_patterns(actual_messages, EXPECTED_PATTERNS)
+                )
+                patterns_match_perfectly = (
+                    len(missing_patterns) == 0 and len(unmatched_messages) == 0
+                )
+
+                _log_pattern_analysis(
+                    actual_messages,
+                    EXPECTED_PATTERNS,
+                    matched_patterns,
+                    missing_patterns,
+                    unmatched_messages,
+                    show_all_messages=patterns_match_perfectly,
+                )
+
+            _assert_log_bytes_within_tolerance(
+                total_log_bytes, EXPECTED_BYTES, ACCEPTABLE_DELTA
+            )
+
+
+@pytest.mark.skipolddriver
+async def test_no_new_warnings_or_errors_on_successful_basic_select(conn_cnx, caplog):
+    """Test that the number of warning/error log entries stays the same during successful basic select operations. Related to: SNOW-2268606"""
+    caplog.set_level(logging.WARNING, "snowflake.connector")
+    baseline_warning_count = 0
+    baseline_error_count = 0
+
+    # Execute basic select operations and check counts remain the same
+    caplog.clear()
+    async with conn_cnx() as conn:
+        async with conn.cursor() as cur:
+            # Execute basic select operations
+            result1 = await (await cur.execute("select 1")).fetchall()
+            assert result1 == [(1,)]
+
+    # Count warning/error log entries after operations
+    test_warning_count = len(
+        [r for r in caplog.records if r.levelno >= logging.WARNING]
+    )
+    test_error_count = len([r for r in caplog.records if r.levelno >= logging.ERROR])
+
+    # Assert counts stay the same (no new warnings or errors)
+    assert test_warning_count == baseline_warning_count, (
+        f"Warning count increased from {baseline_warning_count} to {test_warning_count}. "
+        f"New warnings: {[r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]}"
+    )
+    assert test_error_count == baseline_error_count, (
+        f"Error count increased from {baseline_error_count} to {test_error_count}. "
+        f"New errors: {[r.getMessage() for r in caplog.records if r.levelno >= logging.ERROR]}"
+    )

test/unit/aio/test_session_manager_async.py

@@ -3,11 +3,13 @@
 
 from unittest import mock
 
+import aiohttp
 import pytest
 
 from snowflake.connector.aio._session_manager import (
     AioHttpConfig,
     SessionManager,
+    SnowflakeSSLConnector,
     SnowflakeSSLConnectorFactory,
 )
 from snowflake.connector.constants import OCSPMode
@@ -348,3 +350,87 @@ async def test_pickle_session_manager():
 
     await manager.close()
     await unpickled.close()
+
+
+@pytest.fixture
+def mock_connector_with_factory():
+    """Fixture providing a mock connector factory and connector."""
+    mock_connector_factory = mock.MagicMock()
+    mock_connector = mock.MagicMock()
+    mock_connector_factory.return_value = mock_connector
+    return mock_connector, mock_connector_factory
+
+
+@pytest.mark.parametrize(
+    "ocsp_mode,extra_kwargs,expected_kwargs",
+    [
+        # Test with OCSPMode.FAIL_OPEN + extra kwargs (should all appear)
+        (
+            OCSPMode.FAIL_OPEN,
+            {"timeout": 30, "pool_connections": 10},
+            {
+                "timeout": 30,
+                "pool_connections": 10,
+                "snowflake_ocsp_mode": OCSPMode.FAIL_OPEN,
+            },
+        ),
+        # Test with OCSPMode.FAIL_CLOSED + no extra kwargs
+        (
+            OCSPMode.FAIL_CLOSED,
+            {},
+            {"snowflake_ocsp_mode": OCSPMode.FAIL_CLOSED},
+        ),
+        # Checks that None values also cause kwargs name to occur
+        (
+            None,
+            {},
+            {"snowflake_ocsp_mode": None},
+        ),
+        # Test override by extra kwargs: config has FAIL_OPEN but extra_kwargs override with FAIL_CLOSED
+        (
+            OCSPMode.FAIL_OPEN,
+            {"snowflake_ocsp_mode": OCSPMode.FAIL_CLOSED},
+            {"snowflake_ocsp_mode": OCSPMode.FAIL_CLOSED},
+        ),
+    ],
+)
+async def test_aio_http_config_get_connector_parametrized(
+    mock_connector_with_factory, ocsp_mode, extra_kwargs, expected_kwargs
+):
+    """Test that AioHttpConfig.get_connector properly passes kwargs and snowflake_ocsp_mode to connector factory.
+
+    This mirrors the sync test behavior where:
+    - Config attributes are passed to the factory
+    - Extra kwargs can override config attributes
+    - All resulting attributes appear in the factory call
+    """
+    mock_connector, mock_connector_factory = mock_connector_with_factory
+
+    config = AioHttpConfig(
+        connector_factory=mock_connector_factory, snowflake_ocsp_mode=ocsp_mode
+    )
+    result = config.get_connector(**extra_kwargs)
+
+    # Verify the connector factory was called with correct arguments
+    mock_connector_factory.assert_called_once_with(**expected_kwargs)
+    assert result is mock_connector
+
+
+async def test_aio_http_config_get_connector_with_real_connector_factory():
+    """Test get_connector with the actual SnowflakeSSLConnectorFactory.
+
+    Verifies that with a real factory, we get a real SnowflakeSSLConnector instance
+    with the snowflake_ocsp_mode properly set.
+    """
+    config = AioHttpConfig(
+        connector_factory=SnowflakeSSLConnectorFactory(),
+        snowflake_ocsp_mode=OCSPMode.FAIL_CLOSED,
+    )
+
+    connector = config.get_connector(session_manager=SessionManager())
+
+    # Verify we get a real SnowflakeSSLConnector instance
+    assert isinstance(connector, aiohttp.BaseConnector)
+    assert isinstance(connector, SnowflakeSSLConnector)
+    # Verify snowflake_ocsp_mode was set correctly
+    assert connector._snowflake_ocsp_mode == OCSPMode.FAIL_CLOSED