Add init support for SF log stream

DataDog · Nov 1, 2024 · cf47361 · cf47361
1 parent 4fda68f
commit cf47361
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/aws/logs_monitoring/customized_log_group.py b/aws/logs_monitoring/customized_log_group.py
@@ -26,6 +26,11 @@ def is_lambda_customized_log_group(logstream_name):
     )
 
 
+# For both default and customzied Step Functions log groups, the log_stream starts with "states/"
+def is_step_functions_log_group(logstream_name):
+    return logstream_name.startswith("states/")
+
+
 def get_lambda_function_name_from_logstream_name(logstream_name):
     try:
         # Not match the pattern for customized Lambda log group

diff --git a/aws/logs_monitoring/steps/enums.py b/aws/logs_monitoring/steps/enums.py
@@ -136,7 +136,6 @@ def __init__(self, string, event_source):
     RDS = ("/aws/rds", AwsEventSource.RDS)
     # e.g. sns/us-east-1/123456779121/SnsTopicX
     SNS = ("sns/", AwsEventSource.SNS)
-    STEPFUNCTION = ("/aws/vendedlogs/states", AwsEventSource.STEPFUNCTION)
     TRANSITGATEWAY = ("tgw-attach", AwsEventSource.TRANSITGATEWAY)
 
     def __str__(self):

diff --git a/aws/logs_monitoring/steps/handlers/awslogs_handler.py b/aws/logs_monitoring/steps/handlers/awslogs_handler.py
@@ -13,6 +13,7 @@
 )
 from customized_log_group import (
     is_lambda_customized_log_group,
+    is_step_functions_log_group,
     get_lambda_function_name_from_logstream_name,
 )
 from steps.handlers.aws_attributes import AwsAttributes
@@ -108,6 +109,10 @@ def set_source(self, event):
         if is_lambda_customized_log_group(log_stream):
             self.metadata[DD_SOURCE] = str(AwsEventSource.LAMBDA)
 
+        # Special handling for Step Functions log groups. Regardless of whether the log group is customized, the corresponding log stream starts with 'states/'."
+        if is_step_functions_log_group(log_stream):
+            self.metadata[DD_SOURCE] = str(AwsEventSource.STEPFUNCTION)
+
     def add_cloudwatch_tags_from_cache(self):
         log_group_arn = self.aws_attributes.get_log_group_arn()
         formatted_tags = self.cache_layer.get_cloudwatch_log_group_tags_cache().get(
@@ -159,9 +164,6 @@ def handle_rds_source(self):
             )
 
     def handle_step_function_source(self):
-        if not self.aws_attributes.get_log_stream().startswith("states/"):
-            return
-
         state_machine_arn = self.get_state_machine_arn()
         if not state_machine_arn:
             return