Fix async iterator protocol in _SAConnectionContextManager

CHANGES.txt

@@ -19,6 +19,7 @@ To be included in 1.0.0 (unreleased)
 * Align % formatting in Cursor.executemany() with Cursor.execute(), literal % now need to be doubled in Cursor.executemany() #714
 * Fixed unlimited Pool size not working, this is now working as documented by passing maxsize=0 to create_pool #119
 * Added Pool.closed property as present in aiopg #463
+* Fixed SQLAlchemy connection context iterator #410
 
 
 0.0.22 (2021-11-14)

aiomysql/utils.py

@@ -96,9 +96,19 @@ async def __aexit__(self, exc_type, exc, tb):
 
 
 class _SAConnectionContextManager(_ContextManager):
-    async def __aiter__(self):
-        result = await self._coro
-        return result
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._obj is None:
+            self._obj = await self._coro
+
+        try:
+            return await self._obj.__anext__()
+        except StopAsyncIteration:
+            await self._obj.close()
+            self._obj = None
+            raise
 
 
 class _TransactionContextManager(_ContextManager):

tests/sa/test_sa_connection.py

@@ -455,3 +455,14 @@ async def test_create_table(sa_connect):
 
     res = await conn.execute("SELECT * FROM sa_tbl")
     assert 0 == len(await res.fetchall())
+
+
+@pytest.mark.run_loop
+async def test_async_iter(sa_connect):
+    conn = await sa_connect()
+    await conn.execute(tbl.insert().values(name="second"))
+
+    ret = []
+    async for row in conn.execute(tbl.select()):
+        ret.append(row)
+    assert [(1, "first"), (2, "second")] == ret