Enable step function tracing at forwarder level (#831)

* Initial logic to add fetch and add the lambda tag

* Set tag in env var to avoid fetching everytime

* Moved tag to cloudformation param

* Fixed readme, moved param under advanced param

* Removed context passing

* Removed unused import

* Fixed template condition

* Update readme language for consistency

* Moved tagging outside of just lambda

* black formatting

* Updated language to specify all step functions

* Updated language to call it tracing instead of l2t

* Moved step functions specific logic to handle_step_function_source()

* Updated tests to check for new tag

* renamed to 'Step Functions', missing 's' before

Author: avedmala

aws/logs_monitoring/README.md

@@ -148,7 +148,7 @@ Starting version 3.107.0 a new feature is added to enable Lambda function to sto
 
 ### Upgrade an older version to +3.106.0
 
-Starting version 3.106.0 Lambda function has been updated to add a prefix to cache filenames stored in the S3 bucket configured in `DD_S3_BUCKET_NAME`. This allows to use the same bucket to store cache files from several functions.  
+Starting version 3.106.0 Lambda function has been updated to add a prefix to cache filenames stored in the S3 bucket configured in `DD_S3_BUCKET_NAME`. This allows to use the same bucket to store cache files from several functions.
 Additionally, starting this version, the forwarder will attach custom S3 bucket tags by default to all logs exported to S3. For example, if a service is configured to send logs to a destiantion S3 bucket, the forwarder will add the bucket's tags to the logs while pulling and forwarding the logs.
 
 ### Upgrade an older version to +3.99.0
@@ -388,15 +388,6 @@ SSL encrypted TCP connection, set this parameter to true.
 `DdForwardLog`
 : Set to false to disable log forwarding, while continuing to forward other observability data, such as metrics and traces from Lambda functions.
 
-`DdFetchLambdaTags`
-: Let the Forwarder fetch Lambda tags using GetResources API calls and apply them to logs, metrics, and traces. If set to true, permission `tag:GetResources` will be automatically added to the Lambda execution IAM role.
-
-`DdFetchLogGroupTags`
-: Let the forwarder fetch Log Group tags using ListTagsLogGroup and apply them to logs, metrics, and traces. If set to true, permission `logs:ListTagsLogGroup` will be automatically added to the Lambda execution IAM role.
-
-`DdFetchStepFunctionsTags`
-: Let the Forwarder fetch Step Functions tags using GetResources API calls and apply them to logs and traces (if Step Functions tracing is enabled). If set to true, permission `tag:GetResources` will be automatically added to the Lambda execution IAM role.
-
 ### Log scrubbing (optional)
 
 `RedactIp`
@@ -433,6 +424,18 @@ To test different patterns against your logs, turn on [debug logs](#troubleshoot
 
 ### Advanced (optional)
 
+`DdFetchLambdaTags`
+: Let the Forwarder fetch Lambda tags using GetResources API calls and apply them to logs, metrics, and traces. If set to true, permission `tag:GetResources` will be automatically added to the Lambda execution IAM role.
+
+`DdFetchLogGroupTags`
+: Let the forwarder fetch Log Group tags using ListTagsLogGroup and apply them to logs, metrics, and traces. If set to true, permission `logs:ListTagsLogGroup` will be automatically added to the Lambda execution IAM role.
+
+`DdFetchStepFunctionsTags`
+: Let the Forwarder fetch Step Functions tags using GetResources API calls and apply them to logs and traces (if Step Functions tracing is enabled). If set to true, permission `tag:GetResources` will be automatically added to the Lambda execution IAM role.
+
+`DdStepFunctionTraceEnabled`
+: Set to true to enable tracing for all Step Functions.
+
 `SourceZipUrl`
 : Do not change unless you know what you are doing. Override the default location of the function source code.
 

aws/logs_monitoring/steps/handlers/awslogs_handler.py

@@ -184,6 +184,12 @@ def handle_step_function_source(self):
                 + ",".join(formatted_stepfunctions_tags)
             )
 
+        if os.environ.get("DD_STEP_FUNCTIONS_TRACE_ENABLED", "false").lower() == "true":
+            self.metadata[DD_CUSTOM_TAGS] = ",".join(
+                [self.metadata.get(DD_CUSTOM_TAGS, [])]
+                + ["dd_step_functions_trace_enabled:true"]
+            )
+
     def handle_verified_access_source(self):
         try:
             message = json.loads(self.aws_attributes.get_log_events()[0].get("message"))

aws/logs_monitoring/template.yaml

@@ -161,6 +161,13 @@ Parameters:
       - true
       - false
     Description: Set to false to disable log forwarding, while continuing to forward other observability data, such as metrics and traces from Lambda functions.
+  DdStepFunctionsTraceEnabled:
+    Type: String
+    Default: false
+    AllowedValues:
+      - true
+      - false
+    Description: Set to true to enable tracing for all Step Functions.
   DdUseCompression:
     Type: String
     Default: true
@@ -391,6 +398,10 @@ Conditions:
     Fn::Equals:
       - Ref: DdForwardLog
       - false
+  SetDdStepFunctionsTraceEnabled:
+    Fn::Equals:
+      - Ref: DdStepFunctionsTraceEnabled
+      - true
   SetDdUseCompression:
     Fn::Equals:
       - Ref: DdUseCompression
@@ -615,6 +626,11 @@ Resources:
               - SetDdForwardLog
               - Ref: DdForwardLog
               - Ref: AWS::NoValue
+          DD_STEP_FUNCTIONS_TRACE_ENABLED:
+            Fn::If:
+              - SetDdStepFunctionsTraceEnabled
+              - Ref: DdStepFunctionsTraceEnabled
+              - Ref: AWS::NoValue
           DD_USE_COMPRESSION:
             Fn::If:
               - SetDdUseCompression
@@ -1148,6 +1164,7 @@ Metadata:
           - DdFetchLambdaTags
           - DdFetchLogGroupTags
           - DdFetchStepFunctionsTags
+          - DdStepFunctionsTraceEnabled
           - TagsCacheTTLSeconds
           - SourceZipUrl
           - InstallAsLayer

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",
+    "ddtags": "env:dev,test_tag_key:test_tag_value,dd_step_functions_trace_enabled:true",
     "host": "/aws/vendedlogs/states/logs-to-traces-sequential-Logs",
     "service": "stepfunction"
 }

aws/logs_monitoring/tests/test_awslogs_handler.py

@@ -71,6 +71,7 @@ def test_awslogs_handler_rds_postgresql(self, mock_cache_init):
 
     @patch("caching.cloudwatch_log_group_cache.CloudwatchLogGroupTagsCache.__init__")
     @patch("caching.cloudwatch_log_group_cache.send_forwarder_internal_metrics")
+    @patch.dict("os.environ", {"DD_STEP_FUNCTIONS_TRACE_ENABLED": "true"})
     def test_awslogs_handler_step_functions_tags_added_properly(
         self,
         mock_forward_metrics,