Review comments. Move types to types.py, move some functions to util.py. Tested and verified docs.

Author: OlegZv

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

@@ -126,187 +126,61 @@ def response_hook(span, instance, response):
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.redis.package import _instruments
 from opentelemetry.instrumentation.redis.util import (
-    _extract_conn_attributes,
+    _add_create_attributes,
+    _add_search_attributes,
+    _build_span_meta_data_for_pipeline,
+    _build_span_name,
     _format_command_args,
-    _set_span_attribute_if_value,
-    _value_or_none,
+    _set_connection_attributes,
 )
 from opentelemetry.instrumentation.redis.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import (
-    Span,
     StatusCode,
     Tracer,
     TracerProvider,
     get_tracer,
 )
 
 if TYPE_CHECKING:
-    from typing import Awaitable, TypeVar
+    from typing import Awaitable
 
     import redis.asyncio.client
     import redis.asyncio.cluster
     import redis.client
     import redis.cluster
     import redis.connection
 
-    RequestHook = Callable[
-        [Span, redis.connection.Connection, list[Any], dict[str, Any]], None
-    ]
-    ResponseHook = Callable[[Span, redis.connection.Connection, Any], None]
-
-    AsyncPipelineInstance = TypeVar(
-        "AsyncPipelineInstance",
-        redis.asyncio.client.Pipeline,
-        redis.asyncio.cluster.ClusterPipeline,
-    )
-    AsyncRedisInstance = TypeVar(
-        "AsyncRedisInstance", redis.asyncio.Redis, redis.asyncio.RedisCluster
-    )
-    PipelineInstance = TypeVar(
-        "PipelineInstance",
-        redis.client.Pipeline,
-        redis.cluster.ClusterPipeline,
-    )
-    RedisInstance = TypeVar(
-        "RedisInstance", redis.client.Redis, redis.cluster.RedisCluster
+    from opentelemetry.instrumentation.redis.types import (
+        AsyncPipelineInstance,
+        AsyncRedisInstance,
+        PipelineInstance,
+        R,
+        RedisInstance,
+        RequestHook,
+        ResponseHook,
     )
-    R = TypeVar("R")
 
 
-_DEFAULT_SERVICE = "redis"
 _logger = logging.getLogger(__name__)
 
 _REDIS_ASYNCIO_VERSION = (4, 2, 0)
 _REDIS_CLUSTER_VERSION = (4, 1, 0)
 _REDIS_ASYNCIO_CLUSTER_VERSION = (4, 3, 2)
 
-_FIELD_TYPES = ["NUMERIC", "TEXT", "GEO", "TAG", "VECTOR"]
 
 _CLIENT_ASYNCIO_SUPPORT = redis.VERSION >= _REDIS_ASYNCIO_VERSION
 _CLIENT_ASYNCIO_CLUSTER_SUPPORT = (
     redis.VERSION >= _REDIS_ASYNCIO_CLUSTER_VERSION
 )
 _CLIENT_CLUSTER_SUPPORT = redis.VERSION >= _REDIS_CLUSTER_VERSION
-_CLIENT_BEFORE_3_0_0 = redis.VERSION < (3, 0, 0)
+_CLIENT_BEFORE_V3 = redis.VERSION < (3, 0, 0)
 
 if _CLIENT_ASYNCIO_SUPPORT:
     import redis.asyncio
 
-INSTRUMENTATION_ATTR = "_is_instrumented_by_opentelemetry"
-
-
-def _set_connection_attributes(
-    span: Span, conn: RedisInstance | AsyncRedisInstance
-) -> None:
-    if not span.is_recording() or not hasattr(conn, "connection_pool"):
-        return
-    for key, value in _extract_conn_attributes(
-        conn.connection_pool.connection_kwargs
-    ).items():
-        span.set_attribute(key, value)
-
-
-def _build_span_name(
-    instance: RedisInstance | AsyncRedisInstance, cmd_args: tuple[Any, ...]
-) -> str:
-    if len(cmd_args) > 0 and cmd_args[0]:
-        if cmd_args[0] == "FT.SEARCH":
-            name = "redis.search"
-        elif cmd_args[0] == "FT.CREATE":
-            name = "redis.create_index"
-        else:
-            name = cmd_args[0]
-    else:
-        name = instance.connection_pool.connection_kwargs.get("db", 0)
-    return name
-
-
-def _add_create_attributes(span: Span, args: tuple[Any, ...]):
-    _set_span_attribute_if_value(
-        span, "redis.create_index.index", _value_or_none(args, 1)
-    )
-    # According to: https://github.com/redis/redis-py/blob/master/redis/commands/search/commands.py#L155 schema is last argument for execute command
-    try:
-        schema_index = args.index("SCHEMA")
-    except ValueError:
-        return
-    schema = args[schema_index:]
-    field_attribute = ""
-    # Schema in format:
-    # [first_field_name, first_field_type, first_field_some_attribute1, first_field_some_attribute2, second_field_name, ...]
-    field_attribute = "".join(
-        f"Field(name: {schema[index - 1]}, type: {schema[index]});"
-        for index in range(1, len(schema))
-        if schema[index] in _FIELD_TYPES
-    )
-    _set_span_attribute_if_value(
-        span,
-        "redis.create_index.fields",
-        field_attribute,
-    )
-
-
-def _add_search_attributes(span: Span, response, args):
-    _set_span_attribute_if_value(
-        span, "redis.search.index", _value_or_none(args, 1)
-    )
-    _set_span_attribute_if_value(
-        span, "redis.search.query", _value_or_none(args, 2)
-    )
-    # Parse response from search
-    # https://redis.io/docs/latest/commands/ft.search/
-    # Response in format:
-    # [number_of_returned_documents, index_of_first_returned_doc, first_doc(as a list), index_of_second_returned_doc, second_doc(as a list) ...]
-    # Returned documents in array format:
-    # [first_field_name, first_field_value, second_field_name, second_field_value ...]
-    number_of_returned_documents = _value_or_none(response, 0)
-    _set_span_attribute_if_value(
-        span, "redis.search.total", number_of_returned_documents
-    )
-    if "NOCONTENT" in args or not number_of_returned_documents:
-        return
-    for document_number in range(number_of_returned_documents):
-        document_index = _value_or_none(response, 1 + 2 * document_number)
-        if document_index:
-            document = response[2 + 2 * document_number]
-            for attribute_name_index in range(0, len(document), 2):
-                _set_span_attribute_if_value(
-                    span,
-                    f"redis.search.xdoc_{document_index}.{document[attribute_name_index]}",
-                    document[attribute_name_index + 1],
-                )
-
-
-def _build_span_meta_data_for_pipeline(
-    instance: PipelineInstance | AsyncPipelineInstance,
-) -> tuple[list[Any], str, str]:
-    try:
-        command_stack = (
-            instance.command_stack
-            if hasattr(instance, "command_stack")
-            else instance._command_stack
-        )
-
-        cmds = [
-            _format_command_args(c.args if hasattr(c, "args") else c[0])
-            for c in command_stack
-        ]
-        resource = "\n".join(cmds)
-
-        span_name = " ".join(
-            [
-                (c.args[0] if hasattr(c, "args") else c[0][0])
-                for c in command_stack
-            ]
-        )
-    except (AttributeError, IndexError):
-        command_stack = []
-        resource = ""
-        span_name = ""
-
-    return command_stack, resource, span_name
+_INSTRUMENTATION_ATTR = "_is_instrumented_by_opentelemetry"
 
 
 def _traced_execute_factory(
@@ -479,8 +353,8 @@ def _instrument(
     _traced_execute_pipeline = _traced_execute_pipeline_factory(
         tracer, request_hook, response_hook
     )
-    pipeline_class = "BasePipeline" if _CLIENT_BEFORE_3_0_0 else "Pipeline"
-    redis_class = "StrictRedis" if _CLIENT_BEFORE_3_0_0 else "Redis"
+    pipeline_class = "BasePipeline" if _CLIENT_BEFORE_V3 else "Pipeline"
+    redis_class = "StrictRedis" if _CLIENT_BEFORE_V3 else "Redis"
 
     wrap_function_wrapper(
         "redis", f"{redis_class}.execute_command", _traced_execute_command
@@ -677,7 +551,7 @@ def _instrument(self, **kwargs: Any):
         )
 
     def _uninstrument(self, **kwargs: Any):
-        if _CLIENT_BEFORE_3_0_0:
+        if _CLIENT_BEFORE_V3:
             unwrap(redis.StrictRedis, "execute_command")
             unwrap(redis.StrictRedis, "pipeline")
             unwrap(redis.Redis, "pipeline")
@@ -739,16 +613,16 @@ def instrument_client(
 
                 The ``args`` represents the response.
         """
-        if not hasattr(client, INSTRUMENTATION_ATTR):
-            setattr(client, INSTRUMENTATION_ATTR, False)
-        if not getattr(client, INSTRUMENTATION_ATTR):
+        if not hasattr(client, _INSTRUMENTATION_ATTR):
+            setattr(client, _INSTRUMENTATION_ATTR, False)
+        if not getattr(client, _INSTRUMENTATION_ATTR):
             _instrument_client(
                 client,
                 RedisInstrumentor._get_tracer(tracer_provider=tracer_provider),
                 request_hook=request_hook,
                 response_hook=response_hook,
             )
-            setattr(client, INSTRUMENTATION_ATTR, True)
+            setattr(client, _INSTRUMENTATION_ATTR, True)
         else:
             _logger.warning(
                 "Attempting to instrument Redis connection while already instrumented"
@@ -767,7 +641,7 @@ def uninstrument_client(
         Args:
             client: The redis client
         """
-        if getattr(client, INSTRUMENTATION_ATTR):
+        if getattr(client, _INSTRUMENTATION_ATTR):
             # for all clients we need to unwrap execute_command and pipeline functions
             unwrap(client, "execute_command")
             # the method was creating a pipeline and wrapping the functions of the

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

@@ -0,0 +1,32 @@
+from typing import Any, Callable, TypeVar
+
+import redis.asyncio.client
+import redis.asyncio.cluster
+import redis.client
+import redis.cluster
+import redis.connection
+
+from opentelemetry.trace import Span
+
+RequestHook = Callable[
+    [Span, redis.connection.Connection, list[Any], dict[str, Any]], None
+]
+ResponseHook = Callable[[Span, redis.connection.Connection, Any], None]
+
+AsyncPipelineInstance = TypeVar(
+    "AsyncPipelineInstance",
+    redis.asyncio.client.Pipeline,
+    redis.asyncio.cluster.ClusterPipeline,
+)
+AsyncRedisInstance = TypeVar(
+    "AsyncRedisInstance", redis.asyncio.Redis, redis.asyncio.RedisCluster
+)
+PipelineInstance = TypeVar(
+    "PipelineInstance",
+    redis.client.Pipeline,
+    redis.cluster.ClusterPipeline,
+)
+RedisInstance = TypeVar(
+    "RedisInstance", redis.client.Redis, redis.cluster.RedisCluster
+)
+R = TypeVar("R")