From e95de3d99b828b82424ffd6b21f317099fe52ba5 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <tazarov@users.noreply.github.com>
Date: Tue, 24 Oct 2023 03:00:26 +0300
Subject: [PATCH] [BUG]: Added return of the original function result when
 tracer not defined (#1268)

-
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- When (for some reason) the tracer is not defined, we seem to be
returning None, added original function return for completeness
- Added _transform_attributes method as this seemed to be throwing an
exception

## Test plan
*How are these changes tested?*

Manual testing with Zipkin/Jaeger

## Documentation Changes
N/A
---
 chromadb/telemetry/opentelemetry/__init__.py | 46 ++++++++++++++++----
 1 file changed, 37 insertions(+), 9 deletions(-)

diff --git a/chromadb/telemetry/opentelemetry/__init__.py b/chromadb/telemetry/opentelemetry/__init__.py
index a713ed11ab4..0160c28183d 100644
--- a/chromadb/telemetry/opentelemetry/__init__.py
+++ b/chromadb/telemetry/opentelemetry/__init__.py
@@ -1,6 +1,6 @@
 from functools import wraps
 from enum import Enum
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 from opentelemetry import trace
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
@@ -72,8 +72,10 @@ def otel_init(
     Parameters match the environment variables which configure OTel as documented
     at https://docs.trychroma.com/observability.
     - otel_service_name: The name of the service for OTel tagging and aggregation.
-    - otel_collection_endpoint: The endpoint to which OTel spans are sent (e.g. api.honeycomb.com).
-    - otel_collection_headers: The headers to send with OTel spans (e.g. {"x-honeycomb-team": "abc123"}).
+    - otel_collection_endpoint: The endpoint to which OTel spans are sent
+        (e.g. api.honeycomb.com).
+    - otel_collection_headers: The headers to send with OTel spans
+        (e.g. {"x-honeycomb-team": "abc123"}).
     - otel_granularity: The granularity of the spans to emit.
     """
     if otel_granularity == OpenTelemetryGranularity.NONE:
@@ -99,18 +101,32 @@ def otel_init(
 def trace_method(
     trace_name: str,
     trace_granularity: OpenTelemetryGranularity,
-    attributes: Dict[str, Union[str, bool, float, int]] = {},
+    attributes: Optional[
+        Dict[
+            str,
+            Union[
+                str,
+                bool,
+                float,
+                int,
+                Sequence[str],
+                Sequence[bool],
+                Sequence[float],
+                Sequence[int],
+            ],
+        ]
+    ] = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """A decorator that traces a method."""
 
     def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(f)
         def wrapper(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
-            global tracer, granularity, _transform_attributes
+            global tracer, granularity
             if trace_granularity < granularity:
                 return f(*args, **kwargs)
             if not tracer:
-                return
+                return f(*args, **kwargs)
             with tracer.start_as_current_span(trace_name, attributes=attributes):
                 return f(*args, **kwargs)
 
@@ -120,13 +136,25 @@ def wrapper(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
 
 
 def add_attributes_to_current_span(
-    attributes: Dict[str, Union[str, bool, float, int]]
+    attributes: Dict[
+        str,
+        Union[
+            str,
+            bool,
+            float,
+            int,
+            Sequence[str],
+            Sequence[bool],
+            Sequence[float],
+            Sequence[int],
+        ],
+    ]
 ) -> None:
     """Add attributes to the current span."""
-    global tracer, granularity, _transform_attributes
+    global tracer, granularity
     if granularity == OpenTelemetryGranularity.NONE:
         return
     if not tracer:
         return
     span = trace.get_current_span()
-    span.set_attributes(_transform_attributes(attributes))  # type: ignore
+    span.set_attributes(attributes)