fix: flatten the map for span attributes

Signed-off-by: Shivanshu Raj Shrivastava <[email protected]>

Author: shivanshuraj1333

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py

@@ -55,6 +55,17 @@
 )
 
 
+def flatten_dict(d, parent_key="", sep="."):
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
 # pylint:disable=too-many-branches
 def set_attributes_from_context(span, context):
     """Helper to extract meta values from a Celery Context"""
@@ -95,6 +106,14 @@ def set_attributes_from_context(span, context):
                     SpanAttributes.MESSAGING_DESTINATION, routing_key
                 )
 
+            # Flatten the dictionary
+            if value:
+                flattened = flatten_dict(value)
+                for k, v in flattened.items():
+                    if v is not None and v != "":
+                        span.set_attribute(f"celery.delivery_info.{k}", v)
+            continue
+
         elif key == "id":
             attribute_name = SpanAttributes.MESSAGING_MESSAGE_ID
 
@@ -115,6 +134,15 @@ def set_attributes_from_context(span, context):
                     value = "topic"
                     break
 
+        # If the value is a dictionary, flatten it
+        if isinstance(value, dict):
+            flattened = flatten_dict(value)
+            for nested_key, nested_value in flattened.items():
+                if nested_value is not None and nested_value != "":
+                    nested_attr_name = f"celery.{key}.{nested_key}"
+                    span.set_attribute(nested_attr_name, nested_value)
+            continue
+
         # set attribute name if not set specially for a key
         if attribute_name is None:
             attribute_name = f"celery.{key}"

instrumentation/opentelemetry-instrumentation-celery/tests/test_utils.py

@@ -58,9 +58,7 @@ def test_set_attributes_from_context(self):
             span.attributes.get(SpanAttributes.MESSAGING_DESTINATION), "celery"
         )
 
-        self.assertEqual(
-            span.attributes["celery.delivery_info"], str({"eager": True})
-        )
+        self.assertEqual(span.attributes["celery.delivery_info.eager"], True)
         self.assertEqual(span.attributes.get("celery.eta"), "soon")
         self.assertEqual(span.attributes.get("celery.expires"), "later")
         self.assertEqual(span.attributes.get("celery.hostname"), "localhost")
@@ -72,6 +70,59 @@ def test_set_attributes_from_context(self):
         )
         self.assertNotIn("custom_meta", span.attributes)
 
+    def test_set_nested_attributes_from_context(self):
+        context = {
+            "correlation_id": "44b7f305",
+            "delivery_info": {
+                "eager": True,
+                "routing_key": "api_gateway",
+                "priority": 0,
+                "redelivered": False,
+            },
+            "eta": {"time": "soon", "date": "today"},
+            "expires": "later",
+            "hostname": "localhost",
+            "id": "44b7f305",
+            "reply_to": "44b7f305",
+            "retries": 4,
+            "timelimit": ("now", "later"),
+            "custom_meta": "custom_value",
+            "routing_key": "celery",
+        }
+
+        span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext))
+        utils.set_attributes_from_context(span, context)
+
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.MESSAGING_MESSAGE_ID),
+            "44b7f305",
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.MESSAGING_CONVERSATION_ID),
+            "44b7f305",
+        )
+        self.assertEqual(
+            span.attributes.get(SpanAttributes.MESSAGING_DESTINATION), "celery"
+        )
+
+        self.assertEqual(span.attributes["celery.delivery_info.eager"], True)
+        self.assertEqual(
+            span.attributes["celery.delivery_info.routing_key"], "api_gateway"
+        )
+        self.assertEqual(span.attributes["celery.delivery_info.priority"], 0)
+        self.assertEqual(
+            span.attributes["celery.delivery_info.redelivered"], False
+        )
+        self.assertEqual(span.attributes.get("celery.eta.time"), "soon")
+        self.assertEqual(span.attributes.get("celery.eta.date"), "today")
+
+        self.assertEqual(span.attributes.get("celery.reply_to"), "44b7f305")
+        self.assertEqual(span.attributes.get("celery.retries"), 4)
+        self.assertEqual(
+            span.attributes.get("celery.timelimit"), ("now", "later")
+        )
+        self.assertNotIn("custom_meta", span.attributes)
+
     def test_set_attributes_not_recording(self):
         # it should extract only relevant keys
         context = {