feat: add support for yaaredis (#2944)

Co-authored-by: Brett Langdon <[email protected]>

Author: Julien Danjou

.circleci/config.yml

@@ -817,6 +817,16 @@ jobs:
           docker_services: 'redis'
           pattern: 'aredis$'
           snapshot: true
+
+  yaaredis:
+    <<: *machine_executor
+    parallelism: 4
+    steps:
+      - run_test:
+          docker_services: 'redis'
+          pattern: 'yaaredis$'
+          snapshot: true
+
   redis:
     <<: *machine_executor
     parallelism: 4
@@ -995,6 +1005,7 @@ requires_tests: &requires_tests
     - pyramid
     - pytest
     - aredis
+    - yaaredis
     - redis
     - rediscluster
     - requests
@@ -1079,6 +1090,7 @@ workflows:
       - pyramid: *requires_base_venvs
       - pytest: *requires_base_venvs
       - aredis: *requires_base_venvs
+      - yaaredis: *requires_base_venvs
       - redis: *requires_base_venvs
       - rediscluster: *requires_base_venvs
       - requests: *requires_base_venvs

ddtrace/_monkey.py

@@ -71,6 +71,7 @@
     "pyodbc": True,
     "fastapi": True,
     "dogpile_cache": True,
+    "yaaredis": True,
 }
 
 _LOCK = threading.Lock()

ddtrace/contrib/yaaredis/__init__.py

@@ -0,0 +1,58 @@
+"""
+The yaaredis integration traces yaaredis requests.
+
+
+Enabling
+~~~~~~~~
+
+The yaaredis integration is enabled automatically when using
+:ref:`ddtrace-run<ddtracerun>` or :ref:`patch_all()<patch_all>`.
+
+Or use :ref:`patch()<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"``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+To configure particular yaaredis instances use the :ref:`Pin<Pin>` API::
+
+    import yaaredis
+    from ddtrace 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")
+"""
+
+from ...utils.importlib import require_modules
+
+
+required_modules = ["yaaredis", "yaaredis.client"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+
+        __all__ = ["patch"]

ddtrace/contrib/yaaredis/patch.py

@@ -0,0 +1,66 @@
+import yaaredis
+
+from ddtrace import config
+from ddtrace.vendor import wrapt
+
+from ...internal.utils.wrappers import unwrap
+from ...pin import Pin
+from ..redis.util import _trace_redis_cmd
+from ..redis.util import _trace_redis_execute_pipeline
+from ..redis.util import format_command_args
+
+
+config._add("yaaredis", dict(_default_service="redis"))
+
+
+def patch():
+    """Patch the instrumented methods"""
+    if getattr(yaaredis, "_datadog_patch", False):
+        return
+    setattr(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):
+        setattr(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 _trace_redis_cmd(pin, config.yaaredis, 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.yaaredis, resource, instance):
+        return await func(*args, **kwargs)

docs/index.rst

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

docs/integrations.rst

@@ -393,6 +393,11 @@ Vertica
 ^^^^^^^
 .. automodule:: ddtrace.contrib.vertica
 
+.. _yaaredis:
+
+yaaredis
+^^^^^^^^
+.. automodule:: ddtrace.contrib.yaaredis
 
 .. _wsgi:
 

docs/spelling_wordlist.txt

@@ -156,3 +156,4 @@ whitelist
 workflow
 wsgi
 xfail
+yaaredis

releasenotes/notes/yaaredis-integration-9cf0d99f0261a209.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add yaaredis ≥ 2.0.0 support.

riotfile.py

@@ -1252,6 +1252,18 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                 ],
             },
         ),
+        Venv(
+            name="yaaredis",
+            pys=select_pys(min_version="3.6"),
+            command="pytest {cmdargs} tests/contrib/yaaredis",
+            pkgs={
+                "pytest-asyncio": latest,
+                "yaaredis": [
+                    "~=2.0.0",
+                    latest,
+                ],
+            },
+        ),
         Venv(
             name="snowflake",
             command="pytest {cmdargs} tests/contrib/snowflake",