[Logs forwarder] Fix step function tags fetch (DataDog#855)

* [Logs forwarder] Fix step function tags fetch

* update tests

Author: ge0Aja

aws/logs_monitoring/caching/step_functions_cache.py

@@ -132,7 +132,7 @@ def _get_state_machine_tags(self, state_machine_arn: str):
         except Exception as e:
             self.logger.exception(f"Failed to get Step Functions tags due to {e}")
 
-        if len(response.get("ResourceTagMappingList", {})) > 0:
+        if response and len(response.get("ResourceTagMappingList", {})) > 0:
             resource_dict = response.get("ResourceTagMappingList")[0]
             for a_tag in resource_dict.get("Tags", []):
                 key = sanitize_aws_tag_string(a_tag["Key"], remove_colons=True)

aws/logs_monitoring/steps/handlers/awslogs_handler.py

@@ -160,18 +160,13 @@ def handle_rds_source(self):
 
     def handle_step_function_source(self):
         if not self.aws_attributes.get_log_stream().startswith("states/"):
-            pass
-        state_machine_arn = ""
+            return
 
-        try:
-            state_machine_arn = self.get_state_machine_arn()
-            if state_machine_arn:  # not empty
-                self.metadata[DD_HOST] = state_machine_arn
-        except Exception as e:
-            logger.debug(
-                "Unable to set stepfunction host or get state_machine_arn: %s" % e
-            )
+        state_machine_arn = self.get_state_machine_arn()
+        if not state_machine_arn:
+            return
 
+        self.metadata[DD_HOST] = state_machine_arn
         formatted_stepfunctions_tags = (
             self.cache_layer.get_step_functions_tags_cache().get(state_machine_arn)
         )
@@ -214,12 +209,15 @@ def process_eks_logs(self):
         # In case the conditions above don't match we maintain eks as the source
 
     def get_state_machine_arn(self):
-        message = json.loads(self.aws_attributes.get_log_events()[0].get("message"))
-        if message.get("execution_arn") is not None:
-            execution_arn = message["execution_arn"]
-            arn_tokens = re.split(r"[:/\\]", execution_arn)
-            arn_tokens[5] = "stateMachine"
-            return ":".join(arn_tokens[:7])
+        try:
+            message = json.loads(self.aws_attributes.get_log_events()[0].get("message"))
+            if message.get("execution_arn") is not None:
+                execution_arn = message["execution_arn"]
+                arn_tokens = re.split(r"[:/\\]", execution_arn)
+                arn_tokens[5] = "stateMachine"
+                return ":".join(arn_tokens[:7])
+        except Exception as e:
+            logger.debug("Unable to get state_machine_arn: %s" % e)
         return ""
 
     # Lambda logs can be from either default or customized log group

aws/logs_monitoring/tests/approved_files/TestAWSLogsHandler.test_awslogs_handler_step_functions_tags_added_properly.approved.json

@@ -8,7 +8,7 @@
             }
         },
         "id": "37199773595581154154810589279545129148442535997644275712",
-        "message": "{\"id\":\"1\",\"type\":\"ExecutionStarted\",\"details\":{\"input\":\"{\"Comment\": \"Insert your JSON here\"}\",\"inputDetails\":{\"truncated\":false},\"roleArn\":\"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03\"},\",previous_event_id\":\"0\",\"event_timestamp\":\"1668095539607\",\"execution_arn\":\"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:d0dbefd8-a0f6-b402-da4c-f4863def7456:7fa0cfbe-be28-4a20-9875-73c37f5dc39e\"}",
+        "message": "{\"id\": \"1\",\"type\": \"ExecutionStarted\",\"details\": {\"input\": \"{}\",\"inputDetails\": {\"truncated\": \"false\"},\"roleArn\": \"arn:aws:iam::12345678910:role/service-role/StepFunctions-test-role-a0iurr4pt\"},\"previous_event_id\": \"0\",\"event_timestamp\": \"1716992192441\",\"execution_arn\": \"arn:aws:states:us-east-1:12345678910:execution:StepFunction1:ccccccc-d1da-4c38-b32c-2b6b07d713fa\",\"redrive_count\": \"0\"}",
         "timestamp": 1668095539607
     }
 ]

aws/logs_monitoring/tests/approved_files/TestAWSLogsHandler.test_awslogs_handler_step_functions_tags_added_properly.metadata.approved.json

@@ -1,6 +1,6 @@
 {
     "ddsource": "stepfunction",
     "ddtags": "env:dev,test_tag_key:test_tag_value,dd_step_functions_trace_enabled:true",
-    "host": "/aws/vendedlogs/states/logs-to-traces-sequential-Logs",
+    "host": "arn:aws:states:us-east-1:12345678910:stateMachine:StepFunction1",
     "service": "stepfunction"
 }

aws/logs_monitoring/tests/test_awslogs_handler.py

@@ -93,7 +93,7 @@ def test_awslogs_handler_step_functions_tags_added_properly(
                                         {
                                             "id": "37199773595581154154810589279545129148442535997644275712",
                                             "timestamp": 1668095539607,
-                                            "message": '{"id":"1","type":"ExecutionStarted","details":{"input":"{"Comment": "Insert your JSON here"}","inputDetails":{"truncated":false},"roleArn":"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03"},",previous_event_id":"0","event_timestamp":"1668095539607","execution_arn":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:d0dbefd8-a0f6-b402-da4c-f4863def7456:7fa0cfbe-be28-4a20-9875-73c37f5dc39e"}',
+                                            "message": '{"id": "1","type": "ExecutionStarted","details": {"input": "{}","inputDetails": {"truncated": "false"},"roleArn": "arn:aws:iam::12345678910:role/service-role/StepFunctions-test-role-a0iurr4pt"},"previous_event_id": "0","event_timestamp": "1716992192441","execution_arn": "arn:aws:states:us-east-1:12345678910:execution:StepFunction1:ccccccc-d1da-4c38-b32c-2b6b07d713fa","redrive_count": "0"}',
                                         }
                                     ],
                                 }

aws/logs_monitoring/tools/integration_tests/snapshots/step_functions_log.json~snapshot

@@ -87,20 +87,6 @@
             ],
             "type": "distribution"
           },
-          {
-            "device": null,
-            "host": null,
-            "interval": 10,
-            "metric": "aws.dd_forwarder.local_step_functions_tags_cache_expired",
-            "points": "<redacted from snapshot>",
-            "tags": [
-              "forwardername:test",
-              "forwarder_memorysize:1536",
-              "forwarder_version:<redacted from snapshot>",
-              "event_type:awslogs"
-            ],
-            "type": "distribution"
-          },
           {
             "device": null,
             "host": null,