Add additional attributes for redis.search methods create_index, search

instrumentation/opentelemetry-instrumentation-redis/src/opentelemetry/instrumentation/redis/__init__.py

@@ -94,14 +94,18 @@ def response_hook(span, instance, response):
 from typing import Any, Collection
 
 import redis
+import redis.commands
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import trace
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.redis.package import _instruments
 from opentelemetry.instrumentation.redis.util import (
+    _args_or_none,
+    _check_skip,
     _extract_conn_attributes,
     _format_command_args,
+    _set_span_attribute,
 )
 from opentelemetry.instrumentation.redis.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
@@ -181,7 +185,8 @@ def _instrument(
     def _traced_execute_command(func, instance, args, kwargs):
         query = _format_command_args(args)
         name = _build_span_name(instance, args)
-
+        if _check_skip(name):
+            return func(*args, **kwargs)
         with tracer.start_as_current_span(
             name, kind=trace.SpanKind.CLIENT
         ) as span:
@@ -202,9 +207,9 @@ def _traced_execute_pipeline(func, instance, args, kwargs):
             resource,
             span_name,
         ) = _build_span_meta_data_for_pipeline(instance)
-
         exception = None
-
+        if _check_skip(span_name):
+            return func(*args, **kwargs)
         with tracer.start_as_current_span(
             span_name, kind=trace.SpanKind.CLIENT
         ) as span:
@@ -230,6 +235,59 @@ def _traced_execute_pipeline(func, instance, args, kwargs):
 
         return response
 
+    def _traced_create_index(func, instance, args, kwargs):
+        span_name = "redis.create_index"
+        with tracer.start_as_current_span(span_name) as span:
+            _set_span_attribute(
+                span,
+                "redis.create_index.fields",
+                kwargs.get("fields").__str__(),
+            )
+            _set_span_attribute(
+                span,
+                "redis.create_index.definition",
+                kwargs.get("definition").__str__(),
+            )
+        response = func(*args, **kwargs)
+        return response
+
+    def _traced_search(func, instance, args, kwargs):
+        span_name = "redis.search"
+        with tracer.start_as_current_span(span_name) as span:
+            query = kwargs.get("query") or _args_or_none(args, 0)
+            _set_span_attribute(
+                span,
+                "redis.commands.search.query",
+                query.query_string(),
+            )
+            response = func(*args, **kwargs)
+            _set_span_attribute(
+                span, "redis.commands.search.total", response.total
+            )
+            _set_span_attribute(
+                span, "redis.commands.search.duration", response.duration
+            )
+            for index, doc in enumerate(response.docs):
+                _set_span_attribute(
+                    span, f"redis.commands.search.xdoc_{index}", doc.__str__()
+                )
+        return response
+
+    def _traced_aggregate(func, instance, args, kwargs):
+        span_name = "redis.aggregate"
+        with tracer.start_as_current_span(span_name) as span:
+            query = kwargs.get("query") or _args_or_none(args, 0)
+            _set_span_attribute(
+                span,
+                "redis.commands.aggregate.query",
+                query._query,
+            )
+            response = func(*args, **kwargs)
+            _set_span_attribute(
+                span, "redis.commands.aggregate.results", str(response.rows)
+            )
+        return response
+
     pipeline_class = (
         "BasePipeline" if redis.VERSION < (3, 0, 0) else "Pipeline"
     )
@@ -248,6 +306,21 @@ def _traced_execute_pipeline(func, instance, args, kwargs):
         f"{pipeline_class}.immediate_execute_command",
         _traced_execute_command,
     )
+    wrap_function_wrapper(
+        "redis.commands.search",
+        "Search.create_index",
+        _traced_create_index,
+    )
+    wrap_function_wrapper(
+        "redis.commands.search",
+        "Search.search",
+        _traced_search,
+    )
+    wrap_function_wrapper(
+        "redis.commands.search",
+        "Search.aggregate",
+        _traced_aggregate,
+    )
     if redis.VERSION >= _REDIS_CLUSTER_VERSION:
         wrap_function_wrapper(
             "redis.cluster",
@@ -371,6 +444,9 @@ def _instrument(self, **kwargs):
         )
 
     def _uninstrument(self, **kwargs):
+        unwrap(redis.commands.search.Search, "create_index")
+        unwrap(redis.commands.search.Search, "search")
+        unwrap(redis.commands.search.Search, "aggregate")
         if redis.VERSION < (3, 0, 0):
             unwrap(redis.StrictRedis, "execute_command")
             unwrap(redis.StrictRedis, "pipeline")

instrumentation/opentelemetry-instrumentation-redis/src/opentelemetry/instrumentation/redis/util.py

@@ -68,3 +68,25 @@ def _format_command_args(args):
         out_str = ""
 
     return out_str
+
+
+def _check_skip(name):
+    skip_list = ["FT.SEARCH", "FT.CREATE", "FT.AGGREGATE"]
+    for method in skip_list:
+        if method in name:
+            return True
+    return False
+
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def _args_or_none(args, n):
+    try:
+        return args[n]
+    except IndexError:
+        return None

tests/opentelemetry-docker-tests/tests/redis/test_redis_functional.py

@@ -18,6 +18,15 @@
 import redis
 import redis.asyncio
 
+from redis.exceptions import ResponseError
+from redis.commands.search.indexDefinition import IndexDefinition, IndexType
+from redis.commands.search.aggregation import AggregateRequest
+from redis.commands.search.query import Query
+from redis.commands.search.field import (
+    TextField,
+    VectorField,
+)
+
 from opentelemetry import trace
 from opentelemetry.instrumentation.redis import RedisInstrumentor
 from opentelemetry.semconv.trace import SpanAttributes
@@ -614,3 +623,72 @@ def test_get(self):
         self.assertEqual(
             span.attributes.get(SpanAttributes.DB_STATEMENT), "GET ?"
         )
+
+
+class TestRedisearchInstrument(TestBase):
+    def setUp(self):
+        super().setUp()
+        self.redis_client = redis.Redis(port=6379)
+        self.redis_client.flushall()
+        self.embedding_dim = 256
+        RedisInstrumentor().instrument(tracer_provider=self.tracer_provider)
+        self.prepare_data()
+        self.create_index()
+
+    def tearDown(self):
+        RedisInstrumentor().uninstrument()
+        super().tearDown()
+
+    def prepare_data(self):
+        try:
+            self.redis_client.ft("idx:test_vss").dropindex(True)
+        except ResponseError:
+            print("No such index")
+        item = {"name": "test",
+                "value": "test_value",
+                "embeddings": [0.1] * 256}
+        pipeline = self.redis_client.pipeline()
+        pipeline.json().set(f"test:001", "$", item)
+        res = pipeline.execute()
+        assert False not in res
+
+    def create_index(self):
+        schema =  ( 
+                    TextField("$.name", no_stem=True, as_name="name"),
+                    TextField("$.value", no_stem=True, as_name="value"),
+                    VectorField("$.embeddings",
+                                "FLAT",
+                                {
+                                    "TYPE": "FLOAT32",
+                                    "DIM": self.embedding_dim,
+                                    "DISTANCE_METRIC": "COSINE",
+                                },
+                                as_name="vector",),
+                    )
+        definition = IndexDefinition(prefix=["test:"], index_type=IndexType.JSON)
+        res = self.redis_client.ft("idx:test_vss").create_index(fields=schema, definition=definition)
+        assert "OK" in str(res)
+
+    def test_redis_create_index(self):
+        spans = self.memory_exporter.get_finished_spans()
+        span = next(span for span in spans if span.name == "redis.create_index")
+        assert "redis.create_index.definition" in span.attributes
+        assert "redis.create_index.fields" in span.attributes
+
+    def test_redis_aggregate(self):
+        query = "*"
+        self.redis_client.ft("idx:test_vss").aggregate(AggregateRequest(query).load())
+        spans = self.memory_exporter.get_finished_spans()
+        span = next(span for span in spans if span.name == "redis.aggregate")
+        assert span.attributes.get("redis.commands.aggregate.query") == query
+        assert "redis.commands.aggregate.results" in span.attributes
+
+    def test_redis_query(self):
+        query = "@name:test"
+        res = self.redis_client.ft("idx:test_vss").search(Query(query))
+
+        spans = self.memory_exporter.get_finished_spans()
+        span = next(span for span in spans if span.name == "redis.search")
+
+        assert span.attributes.get("redis.commands.search.query") == query
+        assert span.attributes.get("redis.commands.search.total") == 1