[async] Approach 1 - Idempotent connection.__aenter__ through checking if closed

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/__init__.py

@@ -49,17 +49,11 @@ def __iter__(self) -> Generator[Any, None, SnowflakeConnection]:
         """Make the wrapper iterable like a coroutine."""
         return self.__await__()
 
-    # TODO: below is okay if we make idempotent __aenter__ of SnowflakeConnection class - so check if connected and do not repeat connecting
-    # async def __aenter__(self) -> SnowflakeConnection:
-    #     """Enable async with connect(...) as conn:"""
-    #     self._conn = await self._coro
-    #     return await self._conn.__aenter__()
-
+    # This approach requires idempotent __aenter__ of SnowflakeConnection class - so check if connected and do not repeat connecting
     async def __aenter__(self) -> SnowflakeConnection:
         """Enable async with connect(...) as conn:"""
         self._conn = await self._coro
-        # Connection is already connected by the coroutine
-        return self._conn
+        return await self._conn.__aenter__()
 
     async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
         """Exit async context manager."""

src/snowflake/connector/aio/_connection.py

@@ -168,11 +168,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     async def __aenter__(self) -> SnowflakeConnection:
         """Context manager."""
-        # Idempotent __Aenter__
-        # if self.is_closed():
-        #     await self.connect()
-        # return self
-        await self.connect()
+        # Idempotent __aenter__ - required to be able to use both:
+        #   - with snowflake.connector.aio.SnowflakeConnection(**k)
+        #   - with snowflake.connector.aio.connect(**k)
+        if self.is_closed():
+            await self.connect()
         return self
 
     async def __aexit__(