import asyncio
from ddtrace.vendor import wrapt

from aiopg.utils import _ContextManager

from .. import dbapi
from ...constants import ANALYTICS_SAMPLE_RATE_KEY
from ...ext import sql, AppTypes
from ...pin import Pin
from ...settings import config


class AIOTracedCursor(wrapt.ObjectProxy):
    """ TracedCursor wraps a psql cursor and traces its queries. """

    def __init__(self, cursor, pin):
        super(AIOTracedCursor, self).__init__(cursor)
        pin.onto(self)
        name = pin.app or 'sql'
        self._datadog_name = '%s.query' % name

    @asyncio.coroutine
    def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
        pin = Pin.get_from(self)
        if not pin or not pin.enabled():
            result = yield from method(*args, **kwargs)
            return result
        service = pin.service

        with pin.tracer.trace(self._datadog_name, service=service,
                              resource=resource) as s:
            s.span_type = sql.TYPE
            s.set_tag(sql.QUERY, resource)
            s.set_tags(pin.tags)
            s.set_tags(extra_tags)

            # set analytics sample rate
            s.set_tag(
                ANALYTICS_SAMPLE_RATE_KEY,
                config.aiopg.get_analytics_sample_rate()
            )

            try:
                result = yield from method(*args, **kwargs)
                return result
            finally:
                s.set_metric('db.rowcount', self.rowcount)

    @asyncio.coroutine
    def executemany(self, query, *args, **kwargs):
        # FIXME[matt] properly handle kwargs here. arg names can be different
        # with different libs.
        result = yield from self._trace_method(
            self.__wrapped__.executemany, query, {'sql.executemany': 'true'},
            query, *args, **kwargs)
        return result

    @asyncio.coroutine
    def execute(self, query, *args, **kwargs):
        result = yield from self._trace_method(
            self.__wrapped__.execute, query, {}, query, *args, **kwargs)
        return result

    @asyncio.coroutine
    def callproc(self, proc, args):
        result = yield from self._trace_method(
            self.__wrapped__.callproc, proc, {}, proc, args)
        return result


class AIOTracedConnection(wrapt.ObjectProxy):
    """ TracedConnection wraps a Connection with tracing code. """

    def __init__(self, conn, pin=None, cursor_cls=AIOTracedCursor):
        super(AIOTracedConnection, self).__init__(conn)
        name = dbapi._get_vendor(conn)
        db_pin = pin or Pin(service=name, app=name, app_type=AppTypes.db)
        db_pin.onto(self)
        # wrapt requires prefix of `_self` for attributes that are only in the
        # proxy (since some of our source objects will use `__slots__`)
        self._self_cursor_cls = cursor_cls

    def cursor(self, *args, **kwargs):
        # unfortunately we also need to patch this method as otherwise "self"
        # ends up being the aiopg connection object
        coro = self._cursor(*args, **kwargs)
        return _ContextManager(coro)

    @asyncio.coroutine
    def _cursor(self, *args, **kwargs):
        cursor = yield from self.__wrapped__._cursor(*args, **kwargs)
        pin = Pin.get_from(self)
        if not pin:
            return cursor
        return self._self_cursor_cls(cursor, pin)

ddtrace/contrib/aiopg/connection.py
import asyncio
from ddtrace.vendor import wrapt

from aiopg.utils import _ContextManager

from .. import dbapi
from ...constants import ANALYTICS_SAMPLE_RATE_KEY
from ...ext import sql, AppTypes
from ...pin import Pin
from ...settings import config


class AIOTracedCursor(wrapt.ObjectProxy):
    """ TracedCursor wraps a psql cursor and traces its queries. """

    def __init__(self, cursor, pin):
        super(AIOTracedCursor, self).__init__(cursor)
        pin.onto(self)
        name = pin.app or 'sql'
        self._datadog_name = '%s.query' % name

    @asyncio.coroutine
    def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
        pin = Pin.get_from(self)
        if not pin or not pin.enabled():
            result = yield from method(*args, **kwargs)
            return result
        service = pin.service

        with pin.tracer.trace(self._datadog_name, service=service,
                              resource=resource) as s:
            s.span_type = sql.TYPE
            s.set_tag(sql.QUERY, resource)
            s.set_tags(pin.tags)
            s.set_tags(extra_tags)

            # set analytics sample rate
            s.set_tag(
                ANALYTICS_SAMPLE_RATE_KEY,
                config.aiopg.get_analytics_sample_rate()
            )

            try:
                result = yield from method(*args, **kwargs)
                return result
            finally:
                s.set_metric('db.rowcount', self.rowcount)

    @asyncio.coroutine
    def executemany(self, query, *args, **kwargs):
        # FIXME[matt] properly handle kwargs here. arg names can be different
        # with different libs.
        result = yield from self._trace_method(
            self.__wrapped__.executemany, query, {'sql.executemany': 'true'},
            query, *args, **kwargs)
        return result

    @asyncio.coroutine
    def execute(self, query, *args, **kwargs):
        result = yield from self._trace_method(
            self.__wrapped__.execute, query, {}, query, *args, **kwargs)
        return result

    @asyncio.coroutine
    def callproc(self, proc, args):
        result = yield from self._trace_method(
            self.__wrapped__.callproc, proc, {}, proc, args)
        return result

    def __aiter__(self):
        return self.__wrapped__.__aiter__()


class AIOTracedConnection(wrapt.ObjectProxy):
    """ TracedConnection wraps a Connection with tracing code. """

    def __init__(self, conn, pin=None, cursor_cls=AIOTracedCursor):
        super(AIOTracedConnection, self).__init__(conn)
        name = dbapi._get_vendor(conn)
        db_pin = pin or Pin(service=name, app=name, app_type=AppTypes.db)
        db_pin.onto(self)
        # wrapt requires prefix of `_self` for attributes that are only in the
        # proxy (since some of our source objects will use `__slots__`)
        self._self_cursor_cls = cursor_cls

    def cursor(self, *args, **kwargs):
        # unfortunately we also need to patch this method as otherwise "self"
        # ends up being the aiopg connection object
        coro = self._cursor(*args, **kwargs)
        return _ContextManager(coro)

    @asyncio.coroutine
    def _cursor(self, *args, **kwargs):
        cursor = yield from self.__wrapped__._cursor(*args, **kwargs)
        pin = Pin.get_from(self)
        if not pin:
            return cursor
        return self._self_cursor_cls(cursor, pin)

ddtrace/contrib/aiopg/connection.py
diff --git a/ddtrace/contrib/aiopg/connection.py b/ddtrace/contrib/aiopg/connection.py
index 55af5e08d59..da21fd94d47 100644
--- a/ddtrace/contrib/aiopg/connection.py
+++ b/ddtrace/contrib/aiopg/connection.py
@@ -67,6 +67,9 @@ def callproc(self, proc, args):
             self.__wrapped__.callproc, proc, {}, proc, args)
         return result
 
+    def __aiter__(self):
+        return self.__wrapped__.__aiter__()
+
 
 class AIOTracedConnection(wrapt.ObjectProxy):
     """ TracedConnection wraps a Connection with tracing code. """
diff --git a/tests/contrib/aiopg/py37/__init__.py b/tests/contrib/aiopg/py37/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/contrib/aiopg/py37/test.py b/tests/contrib/aiopg/py37/test.py
new file mode 100644
index 00000000000..493786b9dde
--- /dev/null
+++ b/tests/contrib/aiopg/py37/test.py
@@ -0,0 +1,52 @@
+# 3p
+import aiopg
+
+# project
+from ddtrace.contrib.aiopg.patch import patch, unpatch
+from ddtrace import Pin
+
+# testing
+from tests.contrib.config import POSTGRES_CONFIG
+from tests.contrib.asyncio.utils import AsyncioTestCase, mark_asyncio
+
+
+TEST_PORT = str(POSTGRES_CONFIG['port'])
+
+
+class AiopgTestCase(AsyncioTestCase):
+    # default service
+    TEST_SERVICE = 'postgres'
+
+    def setUp(self):
+        super().setUp()
+        self._conn = None
+        patch()
+
+    def tearDown(self):
+        super().tearDown()
+        if self._conn and not self._conn.closed:
+            self._conn.close()
+
+        unpatch()
+
+    async def _get_conn_and_tracer(self):
+        conn = self._conn = await aiopg.connect(**POSTGRES_CONFIG)
+        Pin.get_from(conn).clone(tracer=self.tracer).onto(conn)
+
+        return conn, self.tracer
+
+    @mark_asyncio
+    async def test_async_generator(self):
+        conn, tracer = await self._get_conn_and_tracer()
+        cursor = await conn.cursor()
+        q = 'select \'foobarblah\''
+        await cursor.execute(q)
+        rows = []
+        async for row in cursor:
+            rows.append(row)
+
+        assert rows == [('foobarblah',)]
+        spans = tracer.writer.pop()
+        assert len(spans) == 1
+        span = spans[0]
+        assert span.name == 'postgres.query'

'async for' requires an object with __aiter__ method, got AIOTracedCursor
## Problem
Using ddtrace and aiopg, if I do:

```python
await cur.execute(query)
async for value in cur:
    yield value
```
If my connection is not patched, I get:
```
TypeError: 'async for' requires an object with __aiter__ method, got AIOTracedCursor
(...)
  File "path/to/my/file.py", line 241, in get_many
    async for value in cur:
```
(if my connection is not patched, it works)

## Analysis

The cursor class is replaced with `AIOTracedCursor` which inherits `wrapt.ObjectProxy`.

Problem is, while thanks to `ObjectProxy`, `AIOTracedCursor().__aiter__()` would most probably work and return whatever the real proxy would return, this is not enough for Python to accept that the cursor is an iterator.

A small example with simple objects:
```python
class A():
    def iter(self):
        return iter([])

    async def aiter(self):
        return iter([])

    def __getattr__(self, attr):
        if attr.endswith("iter__"):
            return getattr(self, attr.strip("_"))
a = A()
```
We implement `a.__iter__()` and `a.__aiter__()` but Python doesn't see it:
```
In [6]: a.__iter__()                                                                                                                                                                                               
Out[6]: <list_iterator at 0x7fdff00de860>

In [7]: a.__aiter__()                                                                                                                                                                                              
Out[7]: <coroutine object A.aiter at 0x7fdff00ddba0>

In [8]: async for e in a: print(e)                                                                                                                                                                                 
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
cell_name in async-def-wrapper()

TypeError: 'async for' requires an object with __aiter__ method, got A                                                                                                                                                                                                   

In [9]: iter(a)                                                                                                                                                                                                    
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-2b64cb055077> in <module>
----> 1 iter(a)

TypeError: 'A' object is not iterable

```
