Change tag merging behavior between message ddtags and upper level dd…

…tags

Override the service tag if it exists.
Merge upper level ddtags and message ddtags in one string.

aws/logs_monitoring/lambda_function.py

@@ -269,8 +269,6 @@ def add_metadata_to_lambda_log(event):
             tags.append(service_tag)
             event[DD_SERVICE] = service_tag.split(":")[1]
     else:
-        # remove the service tag from the custom lambda tags if it exists
-        # as we don't want to add it again
         custom_lambda_tags = [
             tag for tag in custom_lambda_tags if not tag.startswith("service:")
         ]
@@ -329,51 +327,18 @@ def extract_ddtags_from_message(event):
                     logger.debug(f"Failed to extract ddtags from: {event}")
                 return
 
-        event[DD_CUSTOM_TAGS] = merge_tags(event[DD_CUSTOM_TAGS], extracted_ddtags)
-
         # Extract service tag from message.ddtags if exists
-        if "service:" in extracted_ddtags:
-            event[DD_SERVICE] = next(
-                tag[8:]
-                for tag in extracted_ddtags.split(",")
-                if tag.startswith("service:")
+        if "service" in extracted_ddtags:
+            event[DD_SERVICE] = next(tag[8:] for tag in extracted_ddtags.split(",") if tag.startswith("service:"))
+            event[DD_CUSTOM_TAGS] = ",".join(
+                [
+                    tag
+                    for tag in event[DD_CUSTOM_TAGS].split(",")
+                    if not tag.startswith("service")
+                ]
             )
 
-
-def merge_tags(custom_tags, application_tags):
-    """Merge the custom tags added by the forwarder and the application.
-
-    The custom tags added by the forwarder are added to the top-level `ddtags`
-    field, while the custom tags added by the application are added to the
-    `message.ddtags` field.
-
-    Args:
-        custom_tags (str): the custom tags added by the forwarder
-        application_tags (str): the custom tags added by the application
-
-    Returns:
-        str: the merged custom tags
-    """
-    if not custom_tags:
-        return application_tags
-
-    if not application_tags:
-        return custom_tags
-
-    custom_tags_set = set(custom_tags.split(","))
-    application_tags_set = set(application_tags.split(","))
-    application_tags_keys = [tag.split(":")[0] for tag in application_tags_set]
-
-    for application_tag_key in application_tags_keys:
-        custom_tags_set = set(
-            [
-                tag
-                for tag in custom_tags_set
-                if not tag.startswith(f"{application_tag_key}:")
-            ]
-        )
-
-    return ",".join(sorted(list(custom_tags_set.union(application_tags_set))))
+        event[DD_CUSTOM_TAGS] = f"{event[DD_CUSTOM_TAGS]},{extracted_ddtags}"
 
 
 def extract_host_from_cloudtrails(event):

aws/logs_monitoring/tests/test_lambda_function.py

@@ -35,7 +35,7 @@
     enrich,
     transform,
     split,
-    merge_tags,
+    extract_ddtags_from_message,
 )
 from parsing import parse, parse_event_type
 
@@ -317,12 +317,98 @@ def test_extract_trace_payload_valid_trace(self):
 
 
 class TestMergeMessageTags(unittest.TestCase):
-    def test_merge_custom_and_application_tags(self):
-        message_tags = "key0:value0,key1:value1"
-        custom_tags = "key2:value2,key0:value3"
+    message_tags = (
+        '{"ddtags":"service:my_application_service,custom_tag_1:value1"}'
+    )
+    custom_tags = "custom_tag_2:value2,service:my_custom_service"
+
+    def test_extract_ddtags_from_message_str(self):
+        event = {
+            "message": self.message_tags,
+            "ddtags": self.custom_tags,
+            "service": "my_service",
+        }
+
+        extract_ddtags_from_message(event)
+
+        self.assertEqual(
+            event["ddtags"],
+            "custom_tag_2:value2,service:my_application_service,custom_tag_1:value1",
+        )
+        self.assertEqual(
+            event["service"],
+            "my_application_service",
+        )
+
+    def test_extract_ddtags_from_message_dict(self):
+        loaded_message_tags = json.loads(self.message_tags)
+        event = {
+            "message": loaded_message_tags,
+            "ddtags": self.custom_tags,
+            "service": "my_service",
+        }
+
+        extract_ddtags_from_message(event)
+
+        self.assertEqual(
+            event["ddtags"],
+            "custom_tag_2:value2,service:my_application_service,custom_tag_1:value1",
+        )
+        self.assertEqual(
+            event["service"],
+            "my_application_service",
+        )
+
+    def test_extract_ddtags_from_message_service_tag_setting(self):
+        loaded_message_tags = json.loads(self.message_tags)
+        loaded_message_tags["ddtags"] = ",".join([tag for tag in loaded_message_tags["ddtags"].split(",") if not tag.startswith("service:")])
+        event = {
+            "message": loaded_message_tags,
+            "ddtags": self.custom_tags,
+            "service": "my_custom_service",
+        }
+
+        extract_ddtags_from_message(event)
+
+        self.assertEqual(
+            event["ddtags"],
+            "custom_tag_2:value2,service:my_custom_service,custom_tag_1:value1",
+        )
+        self.assertEqual(
+            event["service"],
+            "my_custom_service",
+        )
+
+    def test_extract_ddtags_from_message_multiple_service_tag_values(self):
+        custom_tags = self.custom_tags + ",service:my_custom_service_2"
+        event = {"message": self.message_tags, "ddtags": custom_tags}
+
+        extract_ddtags_from_message(event)
+
+        self.assertEqual(
+            event["ddtags"],
+            "custom_tag_2:value2,service:my_application_service,custom_tag_1:value1",
+        )
+        self.assertEqual(
+            event["service"],
+            "my_application_service",
+        )
+
+    def test_extract_ddtags_from_message_multiple_values_tag(self):
+        loaded_message_tags = json.loads(self.message_tags)
+        loaded_message_tags["ddtags"] += ",custom_tag_3:value4"
+        custom_tags = self.custom_tags + ",custom_tag_3:value3"
+        event = {"message": loaded_message_tags, "ddtags": custom_tags}
+
+        extract_ddtags_from_message(event)
+
+        self.assertEqual(
+            event["ddtags"],
+            "custom_tag_2:value2,custom_tag_3:value3,service:my_application_service,custom_tag_1:value1,custom_tag_3:value4",
+        )
         self.assertEqual(
-            merge_tags(custom_tags, message_tags),
-            "key0:value0,key1:value1,key2:value2",
+            event["service"],
+            "my_application_service",
         )