fix(psycopg): support async iteration [backport 2.9]

ddtrace/contrib/dbapi_async/__init__.py

@@ -37,6 +37,9 @@ async def __aenter__(self):
 
         return self
 
+    def __aiter__(self):
+        return self.__wrapped__.__aiter__()
+
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         # previous versions of the dbapi didn't support context managers. let's
         # reference the func that would be called to ensure that error

releasenotes/notes/fix-psycopg-async-iteration-d48a62f8aa597457.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    psycopg: Ensures traced async cursors return an asynchronous iterator object.

tests/contrib/dbapi_async/test_dbapi_async.py

@@ -95,6 +95,27 @@ async def test_fetchone_wrapped_is_called_and_returned(self):
         assert "__result__" == await traced_cursor.fetchone("arg_1", kwarg1="kwarg1")
         cursor.fetchone.assert_called_once_with("arg_1", kwarg1="kwarg1")
 
+    @mark_asyncio
+    async def test_cursor_async_connection(self):
+        """Checks whether connection can execute operations with async iteration."""
+
+        def method():
+            pass
+
+        async with TracedAsyncCursor(self.cursor, Pin("dbapi_service", tracer=self.tracer), {}) as cursor:
+            await cursor.execute("""select 'one' as x""")
+            await cursor.execute("""select 'blah'""")
+
+            async for row in cursor:
+                spans = self.get_spans()
+                assert len(spans) == 2
+                assert spans[0].name == "postgres.query"
+                assert spans[0].resource == "select ?"
+                assert spans[0].service == "dbapi_service"
+                assert spans[1].name == "postgres.query"
+                assert spans[1].resource == "select ?"
+                assert spans[1].service == "dbapi_service"
+
     @mark_asyncio
     async def test_fetchall_wrapped_is_called_and_returned(self):
         cursor = self.cursor

tests/contrib/psycopg/test_psycopg_async.py

@@ -355,3 +355,21 @@ async def test_cursor_from_connection_shortcut(self):
         rows = await cur.fetchall()
         assert len(rows) == 1, rows
         assert rows[0][0] == "one"
+
+    async def test_cursor_async_connect_execute(self):
+        """Checks whether connection can execute operations with async iteration."""
+
+        async with psycopg.AsyncConnection.connect(**POSTGRES_CONFIG) as conn:
+            async with conn.cursor() as cur:
+                await cur.execute("""select 'one' as x""")
+                await cur.execute("""select 'blah'""")
+
+                async for row in cur:
+                    spans = self.get_spans()
+                    assert len(spans) == 2
+                    assert spans[0].name == "postgres.query"
+                    assert spans[0].resource == "select ?"
+                    assert spans[0].service == "postgres"
+                    assert spans[1].name == "postgres.query"
+                    assert spans[1].resource == "select ?"
+                    assert spans[1].service == "postgres"