Repository: DataDog/dd-trace-py

Issue:
aiomysql: AttributeError: __aenter__
### Which version of dd-trace-py are you using?

`ddtrace==1.2.0`

### Which version of the libraries are you using?

`aiomysql==0.1.1`

### How can we reproduce your problem?

```python
# repro.py
import asyncio

import aiomysql


loop = asyncio.get_event_loop()


async def go():
    pool = await aiomysql.create_pool(
        host="127.0.0.1",
        port=3306,
        user="test",
        password="test",
        db="test",
        loop=loop,
        autocommit=False,
    )

    async with pool.acquire() as conn:
        await conn.ping(reconnect=True)
        async with conn.cursor() as cur:
            await cur.execute("SELECT 10")
            (r,) = await cur.fetchone()
            return r


loop.run_until_complete(go())
```

```
ddtrace-run python repro.py
```

### What is the result that you get?

```
❯ ddtrace-run python repro.py
repro.py:16: RuntimeWarning: coroutine 'AIOTracedConnection.cursor' was never awaited
  async with conn.cursor() as cur:
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
Traceback (most recent call last):
  File "/Users/brett.langdon/datadog/dd-trace-py/repro.py", line 22, in <module>
    loop.run_until_complete(go())
  File "/Users/brett.langdon/.pyenv/versions/3.9.10/lib/python3.9/asyncio/base_events.py", line 642, in run_until_complete
    return future.result()
  File "/Users/brett.langdon/datadog/dd-trace-py/repro.py", line 16, in go
    async with conn.cursor() as cur:
AttributeError: __aenter__
```

### What is the result that you expected?
No attribute error, and `async with conn.cursor()` works as expected.


Discussion:

Resolving patch:
diff --git a/ddtrace/contrib/aiomysql/patch.py b/ddtrace/contrib/aiomysql/patch.py
--- a/ddtrace/contrib/aiomysql/patch.py
+++ b/ddtrace/contrib/aiomysql/patch.py
@@ -82,6 +82,14 @@ async def execute(self, query, *args, **kwargs):
         result = await self._trace_method(self.__wrapped__.execute, query, {}, query, *args, **kwargs)
         return result
 
+    # Explicitly define `__aenter__` and `__aexit__` since they do not get proxied properly
+    async def __aenter__(self):
+        # The base class just returns `self`, but we want the wrapped cursor so we return ourselves
+        return self
+
+    async def __aexit__(self, *args, **kwargs):
+        return await self.__wrapped__.__aexit__(*args, **kwargs)
+
 
 class AIOTracedConnection(wrapt.ObjectProxy):
     def __init__(self, conn, pin=None, cursor_cls=AIOTracedCursor):
@@ -93,15 +101,36 @@ def __init__(self, conn, pin=None, cursor_cls=AIOTracedCursor):
         # proxy (since some of our source objects will use `__slots__`)
         self._self_cursor_cls = cursor_cls
 
-    async def cursor(self, *args, **kwargs):
-        cursor = await self.__wrapped__.cursor(*args, **kwargs)
+    def cursor(self, *args, **kwargs):
+        ctx_manager = self.__wrapped__.cursor(*args, **kwargs)
         pin = Pin.get_from(self)
         if not pin:
-            return cursor
-        return self._self_cursor_cls(cursor, pin)
-
+            return ctx_manager
+
+        # The result of `cursor()` is an `aiomysql.utils._ContextManager`
+        #   which wraps a coroutine (a future) and adds async context manager
+        #   helper functions to it.
+        # https://github.com/aio-libs/aiomysql/blob/8a32f052a16dc3886af54b98f4d91d95862bfb8e/aiomysql/connection.py#L461
+        # https://github.com/aio-libs/aiomysql/blob/7fa5078da31bbc95f5e32a934a4b2b4207c67ede/aiomysql/utils.py#L30-L79
+        # We cannot swap out the result on the future/context manager so
+        #   instead we have to create a new coroutine that returns our
+        #   wrapped cursor
+        # We also cannot turn `def cursor` into `async def cursor` because
+        #   otherwise we will change the result to be a coroutine instead of
+        #   an `aiomysql.utils._ContextManager` which wraps a coroutine. This
+        #   will cause issues with `async with conn.cursor() as cur:` usage.
+        async def _wrap_cursor():
+            cursor = await ctx_manager
+            return self._self_cursor_cls(cursor, pin)
+
+        return type(ctx_manager)(_wrap_cursor())
+
+    # Explicitly define `__aenter__` and `__aexit__` since they do not get proxied properly
     async def __aenter__(self):
-        return self.__wrapped__.__aenter__()
+        return await self.__wrapped__.__aenter__()
+
+    async def __aexit__(self, *args, **kwargs):
+        return await self.__wrapped__.__aexit__(*args, **kwargs)
 
 
 def patch():


Test patch:
diff --git a/tests/snapshots/tests.contrib.aiomysql.test_aiomysql.test_async_with_usage.json b/tests/snapshots/tests.contrib.aiomysql.test_aiomysql.test_async_with_usage.json
new file mode 100644
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aiomysql.test_aiomysql.test_async_with_usage.json
@@ -0,0 +1,30 @@
+[[
+  {
+    "name": "mysql.query",
+    "service": null,
+    "resource": "SELECT 1",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "sql",
+    "meta": {
+      "db.name": "test",
+      "db.user": "test",
+      "out.host": "127.0.0.1",
+      "runtime-id": "f09816a4d96c4c51abaf8769c10b0a7a",
+      "sql.query": "SELECT 1"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "db.rowcount": 1,
+      "db.rownumber": 0,
+      "out.port": 3306,
+      "system.pid": 89502
+    },
+    "duration": 1370000,
+    "start": 1656453688953207000
+  }]]