Revert "chore(yaaredis): remove yaaredis support" (#12480)

- [x] Reverts #12226

Author: erikayasuda

.riot/requirements/1761702.txt

@@ -0,0 +1,23 @@
+#
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/1761702.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.4
+exceptiongroup==1.2.0
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==23.2
+pluggy==1.3.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+pytest-cov==4.1.0
+pytest-mock==3.12.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+yaaredis==3.0.0

.riot/requirements/1cc7b0e.txt

@@ -0,0 +1,25 @@
+#
+# This file is autogenerated by pip-compile with Python 3.8
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/1cc7b0e.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.4
+exceptiongroup==1.2.0
+hypothesis==6.45.0
+importlib-metadata==7.0.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==23.2
+pluggy==1.3.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+pytest-cov==4.1.0
+pytest-mock==3.12.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+yaaredis==2.0.4
+zipp==3.17.0

.riot/requirements/1f1e9b4.txt

@@ -0,0 +1,25 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/1f1e9b4.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.4
+exceptiongroup==1.2.0
+hypothesis==6.45.0
+importlib-metadata==7.0.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==23.2
+pluggy==1.3.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+pytest-cov==4.1.0
+pytest-mock==3.12.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+yaaredis==3.0.0
+zipp==3.17.0

.riot/requirements/9b8251b.txt

@@ -0,0 +1,25 @@
+#
+# This file is autogenerated by pip-compile with Python 3.8
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/9b8251b.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.4
+exceptiongroup==1.2.0
+hypothesis==6.45.0
+importlib-metadata==7.0.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==23.2
+pluggy==1.3.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+pytest-cov==4.1.0
+pytest-mock==3.12.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+yaaredis==3.0.0
+zipp==3.17.0

.riot/requirements/fda8aa6.txt

@@ -0,0 +1,25 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/fda8aa6.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.4
+exceptiongroup==1.2.0
+hypothesis==6.45.0
+importlib-metadata==7.0.0
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==23.2
+pluggy==1.3.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+pytest-cov==4.1.0
+pytest-mock==3.12.0
+pytest-randomly==3.15.0
+sortedcontainers==2.4.0
+tomli==2.0.1
+yaaredis==2.0.4
+zipp==3.17.0

ddtrace/_monkey.py

@@ -93,6 +93,7 @@
     "pyodbc": True,
     "fastapi": True,
     "dogpile_cache": True,
+    "yaaredis": True,
     "asyncpg": True,
     "aws_lambda": True,  # patch only in AWS Lambda environments
     "azure_functions": True,

ddtrace/contrib/_yaaredis.py

@@ -0,0 +1,67 @@
+"""
+The yaaredis integration traces yaaredis requests.
+
+
+Enabling
+~~~~~~~~
+
+The yaaredis integration is enabled automatically when using
+:ref:`ddtrace-run<ddtracerun>` or :ref:`import ddtrace.auto<ddtraceauto>`.
+
+Or use :func:`patch()<ddtrace.patch>` to manually enable the integration::
+
+    from ddtrace import patch
+    patch(yaaredis=True)
+
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.yaaredis["service"]
+
+   The service name reported by default for yaaredis traces.
+
+   This option can also be set with the ``DD_YAAREDIS_SERVICE`` environment
+   variable.
+
+   Default: ``"redis"``
+
+.. py:data:: ddtrace.config.yaaredis["cmd_max_length"]
+
+   Max allowable size for the yaaredis command span tag.
+   Anything beyond the max length will be replaced with ``"..."``.
+
+   This option can also be set with the ``DD_YAAREDIS_CMD_MAX_LENGTH`` environment
+   variable.
+
+   Default: ``1000``
+
+.. py:data:: ddtrace.config.aredis["resource_only_command"]
+
+   The span resource will only include the command executed. To include all
+   arguments in the span resource, set this value to ``False``.
+
+   This option can also be set with the ``DD_REDIS_RESOURCE_ONLY_COMMAND`` environment
+   variable.
+
+   Default: ``True``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+To configure particular yaaredis instances use the :class:`Pin <ddtrace.trace.Pin>` API::
+
+    import yaaredis
+    from ddtrace.trace import Pin
+
+    client = yaaredis.StrictRedis(host="localhost", port=6379)
+
+    # Override service name for this instance
+    Pin.override(client, service="my-custom-queue")
+
+    # Traces reported for this client will now have "my-custom-queue"
+    # as the service name.
+    async def example():
+        await client.get("my-key")
+"""

ddtrace/contrib/internal/yaaredis/patch.py

@@ -0,0 +1,91 @@
+import os
+
+import wrapt
+import yaaredis
+
+from ddtrace import config
+from ddtrace._trace.utils_redis import _instrument_redis_cmd
+from ddtrace._trace.utils_redis import _instrument_redis_execute_pipeline
+from ddtrace.contrib.internal.redis_utils import _run_redis_command_async
+from ddtrace.internal.schema import schematize_service_name
+from ddtrace.internal.utils.deprecations import DDTraceDeprecationWarning
+from ddtrace.internal.utils.formats import CMD_MAX_LEN
+from ddtrace.internal.utils.formats import asbool
+from ddtrace.internal.utils.formats import stringify_cache_args
+from ddtrace.internal.utils.wrappers import unwrap
+from ddtrace.trace import Pin
+from ddtrace.vendor.debtcollector import deprecate
+
+
+config._add(
+    "yaaredis",
+    dict(
+        _default_service=schematize_service_name("redis"),
+        cmd_max_length=int(os.getenv("DD_YAAREDIS_CMD_MAX_LENGTH", CMD_MAX_LEN)),
+        resource_only_command=asbool(os.getenv("DD_REDIS_RESOURCE_ONLY_COMMAND", True)),
+    ),
+)
+
+
+def get_version():
+    # type: () -> str
+    return getattr(yaaredis, "__version__", "")
+
+
+def patch():
+    """Patch the instrumented methods"""
+    deprecate(
+        prefix="The yaaredis module is deprecated.",
+        message="The yaaredis module is deprecated and will be deleted.",
+        category=DDTraceDeprecationWarning,
+        removal_version="3.0.0",
+    )
+
+    if getattr(yaaredis, "_datadog_patch", False):
+        return
+    yaaredis._datadog_patch = True
+
+    _w = wrapt.wrap_function_wrapper
+
+    _w("yaaredis.client", "StrictRedis.execute_command", traced_execute_command)
+    _w("yaaredis.client", "StrictRedis.pipeline", traced_pipeline)
+    _w("yaaredis.pipeline", "StrictPipeline.execute", traced_execute_pipeline)
+    _w("yaaredis.pipeline", "StrictPipeline.immediate_execute_command", traced_execute_command)
+    Pin().onto(yaaredis.StrictRedis)
+
+
+def unpatch():
+    if getattr(yaaredis, "_datadog_patch", False):
+        yaaredis._datadog_patch = False
+
+        unwrap(yaaredis.client.StrictRedis, "execute_command")
+        unwrap(yaaredis.client.StrictRedis, "pipeline")
+        unwrap(yaaredis.pipeline.StrictPipeline, "execute")
+        unwrap(yaaredis.pipeline.StrictPipeline, "immediate_execute_command")
+
+
+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 _instrument_redis_cmd(pin, config.yaaredis, instance, args) as ctx:
+        return await _run_redis_command_async(ctx=ctx, func=func, args=args, kwargs=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 = [stringify_cache_args(c, cmd_max_len=config.yaaredis.cmd_max_length) for c, _ in instance.command_stack]
+    with _instrument_redis_execute_pipeline(pin, config.yaaredis, cmds, instance):
+        return await func(*args, **kwargs)

ddtrace/settings/_config.py

@@ -194,6 +194,7 @@
         "grpc_client",
         "grpc_aio_client",
         "grpc_aio_server",
+        "yaaredis",
     }
 )
 

docs/index.rst

@@ -182,6 +182,8 @@ contacting support.
 +--------------------------------------------------+---------------+----------------+
 | :ref:`wsgi`                                      | \*            | No             |
 +--------------------------------------------------+---------------+----------------+
+| :ref:`yaaredis`                                  | >= 2.0.0      | Yes            |
++--------------------------------------------------+---------------+----------------+
 
 
 .. [1] Libraries that are automatically instrumented when the