Redis branch

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

@@ -0,0 +1,82 @@
+"""OpenTelemetry Redis instrumentation"""
+
+import logging
+import redis
+import redis.commands.search
+from typing import Collection
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+
+from opentelemetry.trace import get_tracer
+
+from opentelemetry.instrumentation.redis.config import Config
+from opentelemetry.instrumentation.redis.wrapper import _wrap
+from opentelemetry.instrumentation.redis.version import __version__
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("redis >= 4.0.0",)
+
+WRAPPED_METHODS = [
+    {
+        "package": redis,
+        "object": "Redis",
+        "method": "ping",
+        "span_name": "redis.ping"
+    },
+    {
+        "package": redis,
+        "object": "Redis",
+        "method": "get_connection_kwargs",
+        "span_name": "redis.getconnectionkwargs"
+    },
+    {
+        "package": redis.commands.search,
+        "object": "Search",
+        "method": "create_index",
+        "span_name": "create_index",
+    },
+    {
+        "package": redis.commands.search,
+        "object": "Search",
+        "method": "search",
+        "span_name": "search",
+    },
+]
+
+class RedisInstrumentor(BaseInstrumentor):
+    """An instrumentor for Redis client library."""
+
+    def __init__(self, exception_logger=None):
+        print("Redis instrumentor initialized __init__")
+        super().__init__()
+        Config.exception_logger = exception_logger
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        print("In redis _instrument")
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrap_method = wrapped_method.get("method")
+            if getattr(wrap_package, wrap_object, None):
+                wrap_function_wrapper(
+                    wrap_package,
+                    f"{wrap_object}.{wrap_method}",
+                    _wrap(tracer, wrapped_method),
+                )
+
+    def _uninstrument(self, **kwargs):
+        print("In redis _uninstrument")
+        for wrapped_method in WRAPPED_METHODS:
+            wrap_package = wrapped_method.get("package")
+            wrap_object = wrapped_method.get("object")
+            wrapped = getattr(wrap_package, wrap_object, None)
+            if wrapped:
+                unwrap(wrapped, wrapped_method.get("method"))

packages/opentelemetry-instrumentation-redis/opentelemetry/instrumentation/redis/config.py

@@ -0,0 +1,2 @@
+class Config:
+    exception_logger = None

packages/opentelemetry-instrumentation-redis/opentelemetry/instrumentation/redis/utils.py

@@ -0,0 +1,23 @@
+import logging
+from opentelemetry.instrumentation.redis.config import Config
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.warning("Failed to execute %s, error: %s", func.__name__, str(e))
+            if Config.exception_logger:
+                Config.exception_logger(e)
+
+    return wrapper

packages/opentelemetry-instrumentation-redis/opentelemetry/instrumentation/redis/version.py

@@ -0,0 +1 @@
+__version__ = "0.0.0"

packages/opentelemetry-instrumentation-redis/opentelemetry/instrumentation/redis/wrapper.py

@@ -0,0 +1,45 @@
+from opentelemetry import context as context_api
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.instrumentation.utils import (
+    _SUPPRESS_INSTRUMENTATION_KEY,
+)
+
+def _with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+def _set_span_attribute(span, name, value):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+def count_or_none(obj):
+    if obj:
+        return len(obj)
+
+    return None
+
+@_with_tracer_wrapper
+def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
+    print("In redis _wrap")
+    """Instruments and calls every function defined in TO_WRAP."""
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    name = to_wrap.get("span_name")
+    with tracer.start_as_current_span(name) as span:
+        print(f"Name of the span: {name}")
+        span.set_attribute(SpanAttributes.DB_SYSTEM, "redis")
+        span.set_attribute(SpanAttributes.DB_OPERATION, to_wrap.get("method"))
+        _set_span_attribute(span, "db.redis.search.test", 1)
+        return_value = wrapped(*args, **kwargs)
+    print(f"Return value: {return_value}")
+    return return_value

packages/opentelemetry-instrumentation-redis/pyproject.toml

@@ -0,0 +1,47 @@
+[tool.coverage.run]
+branch = true
+source = [ "opentelemetry/instrumentation/redis" ]
+
+[tool.coverage.report]
+exclude_lines = [ "if TYPE_CHECKING:" ]
+show_missing = true
+
+[tool.poetry]
+name = "opentelemetry-instrumentation-redis"
+version = "0.16.6"
+description = "OpenTelemetry Redis instrumentation"
+authors = [
+]
+repository = "https://github.com/traceloop/openllmetry/tree/main/packages/opentelemetry-instrumentation-redis"
+license = "Apache-2.0"
+readme = "README.md"
+
+[[tool.poetry.packages]]
+include = "opentelemetry/instrumentation/redis"
+
+[tool.poetry.dependencies]
+python = ">=3.9,<4"
+opentelemetry-api = "^1.24.0"
+opentelemetry-instrumentation = "^0.45b0"
+opentelemetry-semantic-conventions = "^0.45b0"
+opentelemetry-semantic-conventions-ai = "0.1.1"
+
+[tool.poetry.group.dev.dependencies]
+autopep8 = "2.1.0"
+flake8 = "7.0.0"
+
+[tool.poetry.group.test.dependencies]
+redis = "^4.0.0"
+pytest = "8.1.1"
+pytest-sugar = "1.0.0"
+opentelemetry-sdk = "^1.23.0"
+
+[build-system]
+requires = [ "poetry-core" ]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.extras]
+instruments = ["redis"]
+
+[tool.poetry.plugins."opentelemetry_instrumentor"]
+redis_client = "opentelemetry.instrumentation.redis:RedisInstrumentor"