From 0810a46d4821ba094d42b185aaa64b230fec3934 Mon Sep 17 00:00:00 2001
From: Georgi <georgi.ajaeiya@datadoghq.com>
Date: Tue, 22 Oct 2024 15:55:00 +0200
Subject: [PATCH] [Logs forwarder] Encode  log message only once

---
 aws/logs_monitoring/steps/enrichment.py       |  2 +-
 .../tests/events/cloudwatch_logs_ddtags.json  | 16 ++++++++++++
 .../tests/test_lambda_function.py             | 25 +++++++++++++++++--
 3 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 aws/logs_monitoring/tests/events/cloudwatch_logs_ddtags.json

diff --git a/aws/logs_monitoring/steps/enrichment.py b/aws/logs_monitoring/steps/enrichment.py
index 354f301a3..5abdb55a6 100644
--- a/aws/logs_monitoring/steps/enrichment.py
+++ b/aws/logs_monitoring/steps/enrichment.py
@@ -147,7 +147,7 @@ def extract_ddtags_from_message(event):
             try:
                 message_dict = json.loads(event["message"])
                 extracted_ddtags = message_dict.pop(DD_CUSTOM_TAGS)
-                event["message"] = json.dumps(message_dict)
+                event["message"] = message_dict
             except Exception:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(f"Failed to extract ddtags from: {event}")
diff --git a/aws/logs_monitoring/tests/events/cloudwatch_logs_ddtags.json b/aws/logs_monitoring/tests/events/cloudwatch_logs_ddtags.json
new file mode 100644
index 000000000..cc45eff47
--- /dev/null
+++ b/aws/logs_monitoring/tests/events/cloudwatch_logs_ddtags.json
@@ -0,0 +1,16 @@
+{
+	"messageType": "DATA_MESSAGE",
+	"owner": "601427279990",
+	"logGroup": "/aws/lambda/testing-datadog",
+	"logStream": "2024/10/10/[$LATEST]20bddfd5a2dc4c6b97ac02800eae90d0",
+	"subscriptionFilters": [
+		"testing-datadog"
+	],
+	"logEvents": [
+		{
+			"id": "35311576111948622874033876462979853992919938886093242368",
+			"timestamp": 1583425836114,
+			"message": "{\"status\":\"debug\",\"message\":\"datadog:Patched console output with trace context\",\"ddtags\":\"env:test,service:test-inner-message\"}\n"
+		}
+	]
+}
diff --git a/aws/logs_monitoring/tests/test_lambda_function.py b/aws/logs_monitoring/tests/test_lambda_function.py
index 42c94cd35..3e6598fdb 100644
--- a/aws/logs_monitoring/tests/test_lambda_function.py
+++ b/aws/logs_monitoring/tests/test_lambda_function.py
@@ -242,9 +242,30 @@ def test_overrding_service_tag_from_lambda_cache_when_dd_tags_is_set(
         for log in logs:
             self.assertEqual(log["service"], "lambda_service")
 
-    def _get_input_data(self):
+    @patch("caching.cloudwatch_log_group_cache.CloudwatchLogGroupTagsCache.__init__")
+    def test_overrding_service_tag_from_message_ddtags(self, mock_cache_init):
+        mock_cache_init.return_value = None
+        cache_layer = CacheLayer("")
+        cache_layer._lambda_cache = MagicMock()
+        cache_layer._cloudwatch_log_group_cache = MagicMock()
+        context = Context()
+        input_data = self._get_input_data(path="events/cloudwatch_logs_ddtags.json")
+        event = {
+            "awslogs": {"data": self._create_cloudwatch_log_event_from_data(input_data)}
+        }
+        normalized_events = parse(event, context, cache_layer)
+        enriched_events = enrich(normalized_events, cache_layer)
+        transformed_events = transform(enriched_events)
+
+        _, logs, _ = split(transformed_events)
+        self.assertEqual(len(logs), 1)
+        for log in logs:
+            self.assertEqual(log["service"], "test-inner-message")
+            self.assertTrue(isinstance(log["message"], dict))
+
+    def _get_input_data(self, path="events/cloudwatch_logs.json"):
         my_path = os.path.abspath(os.path.dirname(__file__))
-        path = os.path.join(my_path, "events/cloudwatch_logs.json")
+        path = os.path.join(my_path, path)
 
         with open(
             path,