import sys

import aioredis

from ddtrace import config
from ddtrace.internal.utils.wrappers import unwrap as _u
from ddtrace.pin import Pin
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w

from .. import trace_utils
from ...constants import ANALYTICS_SAMPLE_RATE_KEY
from ...constants import SPAN_MEASURED_KEY
from ...ext import SpanTypes
from ...ext import net
from ...ext import redis as redisx
from ..redis.util import _trace_redis_cmd
from ..redis.util import _trace_redis_execute_pipeline
from ..redis.util import format_command_args


try:
    from aioredis.commands.transaction import _RedisBuffer
except ImportError:
    _RedisBuffer = None

config._add("aioredis", dict(_default_service="redis"))

aioredis_version_str = getattr(aioredis, "__version__", "0.0.0")
aioredis_version = tuple([int(i) for i in aioredis_version_str.split(".")])


def patch():
    if getattr(aioredis, "_datadog_patch", False):
        return
    setattr(aioredis, "_datadog_patch", True)
    pin = Pin()
    if aioredis_version >= (2, 0):
        _w("aioredis.client", "Redis.execute_command", traced_execute_command)
        _w("aioredis.client", "Redis.pipeline", traced_pipeline)
        _w("aioredis.client", "Pipeline.execute", traced_execute_pipeline)
        pin.onto(aioredis.client.Redis)
    else:
        _w("aioredis", "Redis.execute", traced_13_execute_command)
        _w("aioredis", "Redis.pipeline", traced_13_pipeline)
        _w("aioredis.commands.transaction", "Pipeline.execute", traced_13_execute_pipeline)
        pin.onto(aioredis.Redis)


def unpatch():
    if not getattr(aioredis, "_datadog_patch", False):
        return

    setattr(aioredis, "_datadog_patch", False)
    if aioredis_version >= (2, 0):
        _u(aioredis.client.Redis, "execute_command")
        _u(aioredis.client.Redis, "pipeline")
        _u(aioredis.client.Pipeline, "execute")
    else:
        _u(aioredis.Redis, "execute")
        _u(aioredis.Redis, "pipeline")
        _u(aioredis.commands.transaction.Pipeline, "execute")


async def traced_execute_command(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    with _trace_redis_cmd(pin, config.aioredis, instance, args):
        return await func(*args, **kwargs)


async def traced_pipeline(func, instance, args, kwargs):
    pipeline = await func(*args, **kwargs)
    pin = Pin.get_from(instance)
    if pin:
        pin.onto(pipeline)
    return pipeline


async def traced_execute_pipeline(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    cmds = [format_command_args(c) for c, _ in instance.command_stack]
    resource = "\n".join(cmds)
    with _trace_redis_execute_pipeline(pin, config.aioredis, resource, instance):
        return await func(*args, **kwargs)


def traced_13_pipeline(func, instance, args, kwargs):
    pipeline = func(*args, **kwargs)
    pin = Pin.get_from(instance)
    if pin:
        pin.onto(pipeline)
    return pipeline


def traced_13_execute_command(func, instance, args, kwargs):
    # If we have a _RedisBuffer then we are in a pipeline
    if isinstance(instance.connection, _RedisBuffer):
        return func(*args, **kwargs)

    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return func(*args, **kwargs)

    # Don't activate the span since this operation is performed as a future which concludes sometime later on in
    # execution so subsequent operations in the stack are not necessarily semantically related
    # (we don't want this span to be the parent of all other spans created before the future is resolved)
    span = pin.tracer.start_span(
        redisx.CMD, service=trace_utils.ext_service(pin, config.aioredis), span_type=SpanTypes.REDIS, activate=False
    )

    span.set_tag(SPAN_MEASURED_KEY)
    query = format_command_args(args)
    span.resource = query
    span.set_tag(redisx.RAWCMD, query)
    if pin.tags:
        span.set_tags(pin.tags)

    span.set_tags(
        {
            net.TARGET_HOST: instance.address[0],
            net.TARGET_PORT: instance.address[1],
            redisx.DB: instance.db or 0,
        }
    )
    span.set_metric(redisx.ARGS_LEN, len(args))
    # set analytics sample rate if enabled
    span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.aioredis.get_analytics_sample_rate())

    def _finish_span(future):
        try:
            # Accessing the result will raise an exception if:
            #   - The future was cancelled
            #   - There was an error executing the future (`future.exception()`)
            #   - The future is in an invalid state
            future.result()
        except Exception:
            span.set_exc_info(*sys.exc_info())
        finally:
            span.finish()

    task = func(*args, **kwargs)
    task.add_done_callback(_finish_span)
    return task


async def traced_13_execute_pipeline(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    cmds = []
    for _, cmd, cmd_args, _ in instance._pipeline:
        parts = [cmd]
        parts.extend(cmd_args)
        cmds.append(format_command_args(parts))
    resource = "\n".join(cmds)
    with pin.tracer.trace(
        redisx.CMD,
        resource=resource,
        service=trace_utils.ext_service(pin, config.aioredis),
        span_type=SpanTypes.REDIS,
    ) as span:

        span.set_tags(
            {
                net.TARGET_HOST: instance._pool_or_conn.address[0],
                net.TARGET_PORT: instance._pool_or_conn.address[1],
                redisx.DB: instance._pool_or_conn.db or 0,
            }
        )

        span.set_tag(SPAN_MEASURED_KEY)
        span.set_tag(redisx.RAWCMD, resource)
        span.set_metric(redisx.PIPELINE_LEN, len(instance._pipeline))
        # set analytics sample rate if enabled
        span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.aioredis.get_analytics_sample_rate())

        return await func(*args, **kwargs)

ddtrace/contrib/aioredis/patch.py
import sys

import aioredis

from ddtrace import config
from ddtrace.internal.utils.wrappers import unwrap as _u
from ddtrace.pin import Pin
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w

from .. import trace_utils
from ...constants import ANALYTICS_SAMPLE_RATE_KEY
from ...constants import SPAN_MEASURED_KEY
from ...ext import SpanTypes
from ...ext import net
from ...ext import redis as redisx
from ..redis.util import _trace_redis_cmd
from ..redis.util import _trace_redis_execute_pipeline
from ..redis.util import format_command_args


try:
    from aioredis.commands.transaction import _RedisBuffer
except ImportError:
    _RedisBuffer = None

config._add("aioredis", dict(_default_service="redis"))

aioredis_version_str = getattr(aioredis, "__version__", "0.0.0")
aioredis_version = tuple([int(i) for i in aioredis_version_str.split(".")])


def patch():
    if getattr(aioredis, "_datadog_patch", False):
        return
    setattr(aioredis, "_datadog_patch", True)
    pin = Pin()
    if aioredis_version >= (2, 0):
        _w("aioredis.client", "Redis.execute_command", traced_execute_command)
        _w("aioredis.client", "Redis.pipeline", traced_pipeline)
        _w("aioredis.client", "Pipeline.execute", traced_execute_pipeline)
        pin.onto(aioredis.client.Redis)
    else:
        _w("aioredis", "Redis.execute", traced_13_execute_command)
        _w("aioredis", "Redis.pipeline", traced_13_pipeline)
        _w("aioredis.commands.transaction", "Pipeline.execute", traced_13_execute_pipeline)
        pin.onto(aioredis.Redis)


def unpatch():
    if not getattr(aioredis, "_datadog_patch", False):
        return

    setattr(aioredis, "_datadog_patch", False)
    if aioredis_version >= (2, 0):
        _u(aioredis.client.Redis, "execute_command")
        _u(aioredis.client.Redis, "pipeline")
        _u(aioredis.client.Pipeline, "execute")
    else:
        _u(aioredis.Redis, "execute")
        _u(aioredis.Redis, "pipeline")
        _u(aioredis.commands.transaction.Pipeline, "execute")


async def traced_execute_command(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    with _trace_redis_cmd(pin, config.aioredis, instance, args):
        return await func(*args, **kwargs)


def traced_pipeline(func, instance, args, kwargs):
    pipeline = func(*args, **kwargs)
    pin = Pin.get_from(instance)
    if pin:
        pin.onto(pipeline)
    return pipeline


async def traced_execute_pipeline(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    cmds = [format_command_args(c) for c, _ in instance.command_stack]
    resource = "\n".join(cmds)
    with _trace_redis_execute_pipeline(pin, config.aioredis, resource, instance):
        return await func(*args, **kwargs)


def traced_13_pipeline(func, instance, args, kwargs):
    pipeline = func(*args, **kwargs)
    pin = Pin.get_from(instance)
    if pin:
        pin.onto(pipeline)
    return pipeline


def traced_13_execute_command(func, instance, args, kwargs):
    # If we have a _RedisBuffer then we are in a pipeline
    if isinstance(instance.connection, _RedisBuffer):
        return func(*args, **kwargs)

    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return func(*args, **kwargs)

    # Don't activate the span since this operation is performed as a future which concludes sometime later on in
    # execution so subsequent operations in the stack are not necessarily semantically related
    # (we don't want this span to be the parent of all other spans created before the future is resolved)
    span = pin.tracer.start_span(
        redisx.CMD, service=trace_utils.ext_service(pin, config.aioredis), span_type=SpanTypes.REDIS, activate=False
    )

    span.set_tag(SPAN_MEASURED_KEY)
    query = format_command_args(args)
    span.resource = query
    span.set_tag(redisx.RAWCMD, query)
    if pin.tags:
        span.set_tags(pin.tags)

    span.set_tags(
        {
            net.TARGET_HOST: instance.address[0],
            net.TARGET_PORT: instance.address[1],
            redisx.DB: instance.db or 0,
        }
    )
    span.set_metric(redisx.ARGS_LEN, len(args))
    # set analytics sample rate if enabled
    span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.aioredis.get_analytics_sample_rate())

    def _finish_span(future):
        try:
            # Accessing the result will raise an exception if:
            #   - The future was cancelled
            #   - There was an error executing the future (`future.exception()`)
            #   - The future is in an invalid state
            future.result()
        except Exception:
            span.set_exc_info(*sys.exc_info())
        finally:
            span.finish()

    task = func(*args, **kwargs)
    task.add_done_callback(_finish_span)
    return task


async def traced_13_execute_pipeline(func, instance, args, kwargs):
    pin = Pin.get_from(instance)
    if not pin or not pin.enabled():
        return await func(*args, **kwargs)

    cmds = []
    for _, cmd, cmd_args, _ in instance._pipeline:
        parts = [cmd]
        parts.extend(cmd_args)
        cmds.append(format_command_args(parts))
    resource = "\n".join(cmds)
    with pin.tracer.trace(
        redisx.CMD,
        resource=resource,
        service=trace_utils.ext_service(pin, config.aioredis),
        span_type=SpanTypes.REDIS,
    ) as span:

        span.set_tags(
            {
                net.TARGET_HOST: instance._pool_or_conn.address[0],
                net.TARGET_PORT: instance._pool_or_conn.address[1],
                redisx.DB: instance._pool_or_conn.db or 0,
            }
        )

        span.set_tag(SPAN_MEASURED_KEY)
        span.set_tag(redisx.RAWCMD, resource)
        span.set_metric(redisx.PIPELINE_LEN, len(instance._pipeline))
        # set analytics sample rate if enabled
        span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.aioredis.get_analytics_sample_rate())

        return await func(*args, **kwargs)

ddtrace/contrib/aioredis/patch.py
diff --git a/ddtrace/contrib/aioredis/patch.py b/ddtrace/contrib/aioredis/patch.py
index c2752835c9f..286e65c0053 100644
--- a/ddtrace/contrib/aioredis/patch.py
+++ b/ddtrace/contrib/aioredis/patch.py
@@ -70,8 +70,8 @@ async def traced_execute_command(func, instance, args, kwargs):
         return await func(*args, **kwargs)
 
 
-async def traced_pipeline(func, instance, args, kwargs):
-    pipeline = await func(*args, **kwargs)
+def traced_pipeline(func, instance, args, kwargs):
+    pipeline = func(*args, **kwargs)
     pin = Pin.get_from(instance)
     if pin:
         pin.onto(pipeline)
diff --git a/releasenotes/notes/fix-aioredis-async-with-pipeline-805966300810edf8.yaml b/releasenotes/notes/fix-aioredis-async-with-pipeline-805966300810edf8.yaml
new file mode 100644
index 00000000000..e09f30b640a
--- /dev/null
+++ b/releasenotes/notes/fix-aioredis-async-with-pipeline-805966300810edf8.yaml
@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Fixes incompatibility of wrapped aioredis pipelines in ``async with`` statements.
diff --git a/tests/contrib/aioredis/test_aioredis.py b/tests/contrib/aioredis/test_aioredis.py
index 33e736c2345..b62d74f35dd 100644
--- a/tests/contrib/aioredis/test_aioredis.py
+++ b/tests/contrib/aioredis/test_aioredis.py
@@ -153,6 +153,35 @@ async def test_pipeline_traced(redis_client):
     assert response_list[3].decode() == "bar"
 
 
+@pytest.mark.skipif(aioredis_version < (2, 0), reason="only supported in aioredis >= 2.0")
+@pytest.mark.asyncio
+@pytest.mark.snapshot
+async def test_pipeline_traced_context_manager_transaction(redis_client):
+    """
+    Regression test for: https://github.com/DataDog/dd-trace-py/issues/3106
+
+    https://aioredis.readthedocs.io/en/latest/migration/#pipelines-and-transactions-multiexec
+
+    Example::
+
+        async def main():
+            redis = await aioredis.from_url("redis://localhost")
+            async with redis.pipeline(transaction=True) as pipe:
+                ok1, ok2 = await (pipe.set("key1", "value1").set("key2", "value2").execute())
+            assert ok1
+            assert ok2
+    """
+
+    async with redis_client.pipeline(transaction=True) as p:
+        set_1, set_2, get_1, get_2 = await (p.set("blah", "boo").set("foo", "bar").get("blah").get("foo").execute())
+
+    # response from redis.set is OK if successfully pushed
+    assert set_1 is True
+    assert set_2 is True
+    assert get_1.decode() == "boo"
+    assert get_2.decode() == "bar"
+
+
 @pytest.mark.asyncio
 @pytest.mark.snapshot(variants={"": aioredis_version >= (2, 0), "13": aioredis_version < (2, 0)})
 async def test_two_traced_pipelines(redis_client):
diff --git a/tests/snapshots/tests.contrib.aioredis.test_aioredis.test_pipeline_traced_context_manager_transaction.json b/tests/snapshots/tests.contrib.aioredis.test_aioredis.test_pipeline_traced_context_manager_transaction.json
new file mode 100644
index 00000000000..b969f3520fb
--- /dev/null
+++ b/tests/snapshots/tests.contrib.aioredis.test_aioredis.test_pipeline_traced_context_manager_transaction.json
@@ -0,0 +1,28 @@
+[[
+  {
+    "name": "redis.command",
+    "service": "redis",
+    "resource": "SET blah boo\nSET foo bar\nGET blah\nGET foo",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "redis",
+    "meta": {
+      "out.host": "127.0.0.1",
+      "redis.raw_command": "SET blah boo\nSET foo bar\nGET blah\nGET foo",
+      "runtime-id": "b734eb991b1f45f2b063db6d3c5623b9"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "out.port": 6379,
+      "out.redis_db": 0,
+      "redis.pipeline_length": 4,
+      "system.pid": 28312
+    },
+    "duration": 2132000,
+    "start": 1641496497488785000
+  }]]

aioredis integration: Tracing breaks pipeline as context managers
Hello,

looks like the `aioredis` integration breaks the interface of `Pipeline` objects as context managers:

```py
RuntimeWarning: coroutine 'traced_pipeline' was never awaited
```

and

```py
async with redis.pipeline(transaction=True) as pipe:
AttributeError: __aexit__
```

This is bad since the documented of usage is exactly as context managers (see https://aioredis.readthedocs.io/en/latest/migration/#pipelines-and-transactions-multiexec).

The fix for now is to just use pipelines outside of contexts, without relying on them as context managers, but that is less than ideal.

`ddtrace` is the latest version (`0.57.0`).
