[Logs forwarder] Fix step function tags fetch

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