Consistent APM Span tags for AWS SDK Requests 2

lib/datadog/tracing/contrib/aws/ext.rb

@@ -31,6 +31,8 @@ module Ext
           TAG_STREAM_NAME = 'streamname'
           TAG_RULE_NAME = 'rulename'
           TAG_STATE_MACHINE_NAME = 'statemachinename'
+          TAG_STATE_MACHINE_ARN = 'statemachinearn'
+          TAG_STATE_EXECUTION_ARN = 'execution_arn'
           TAG_BUCKET_NAME = 'bucketname'
           PEER_SERVICE_SOURCES = Array[TAG_QUEUE_NAME,
             TAG_TOPIC_NAME,

lib/datadog/tracing/contrib/aws/service/states.rb

@@ -14,23 +14,33 @@ def add_tags(span, params)
               state_machine_name = params[:name]
               state_machine_arn = params[:state_machine_arn]
               execution_arn = params[:execution_arn]
-              state_machine_account_id = ''
+              state_machine_account_id = nil
 
+              if state_machine_arn
+                span.set_tag(Aws::Ext::TAG_STATE_MACHINE_ARN, state_machine_arn)
+                # https://docs.aws.amazon.com/step-functions/latest/apireference/API_StartExecution.html#API_StartExecution_RequestSyntax:~:text=Required%3A%20No-,stateMachineArn,-The%20Amazon%20Resource
+                # arn:<partition>:states:<region>:<account-id>:stateMachine:<myStateMachineName>
+                # arn:<partition>:states:<region>:<account-id>:stateMachine:<myStateMachineName>:10
+                # arn:<partition>:states:<region>:<account-id>:stateMachine:<myStateMachineName:PROD>
+                # 3 patterns to cover and attempt to capture the `myStateMachineName` at index 6 and account_id 4
+                parts = state_machine_arn.split(':')
+
+                state_machine_name ||= parts[6] if state_machine_name.nil?
+                state_machine_account_id = parts[4]
+              end
               if execution_arn
-                # 'arn:aws:states:us-east-1:123456789012:execution:example-state-machine:example-execution'
+                span.set_tag(Aws::Ext::TAG_STATE_EXECUTION_ARN, execution_arn)
+                # express
+                # arn:aws:states:sa-east-1:123456789012:express:targetStateMachineName:1234:5678
+                # standard
+                # arn:aws:states:sa-east-1:123456789012:execution:targetStateMachineName:1234
                 parts = execution_arn.split(':')
-                state_machine_name = parts[-2]
+                state_machine_name ||= parts[6] if state_machine_name.nil?
                 state_machine_account_id = parts[4]
               end
 
-              if state_machine_arn
-                # example statemachinearn: arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine
-                parts = state_machine_arn.split(':')
-                state_machine_name ||= parts[-1]
-                state_machine_account_id = parts[-3]
-              end
-              span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, state_machine_account_id)
-              span.set_tag(Aws::Ext::TAG_STATE_MACHINE_NAME, state_machine_name)
+              span.set_tag(Aws::Ext::TAG_AWS_ACCOUNT, state_machine_account_id) if state_machine_account_id
+              span.set_tag(Aws::Ext::TAG_STATE_MACHINE_NAME, state_machine_name) if state_machine_name
             end
           end
         end

spec/datadog/tracing/contrib/aws/service/states_spec.rb

@@ -12,14 +12,22 @@
   end
 
   context 'with execution_arn provided' do
-    let(:execution_arn) { 'arn:aws:states:us-east-1:123456789012:execution:example-state-machine:example-execution' }
+    let(:execution_arn) { 'arn:aws:states:sa-east-1:123456789012:express:targetStateMachineName:1234:5678' }
     let(:params) { { execution_arn: execution_arn } }
 
-    it 'sets the state_machine_name based on the execution_arn' do
+    it 'sets the state_machine_name, execution_arn, and aws_account based on the execution_arn' do
       step_functions.add_tags(span, params)
       expect(span).to have_received(:set_tag).with(
         Datadog::Tracing::Contrib::Aws::Ext::TAG_STATE_MACHINE_NAME,
-        'example-state-machine'
+        'targetStateMachineName'
+      )
+      expect(span).to have_received(:set_tag).with(
+        Datadog::Tracing::Contrib::Aws::Ext::TAG_STATE_EXECUTION_ARN,
+        'arn:aws:states:sa-east-1:123456789012:express:targetStateMachineName:1234:5678'
+      )
+      expect(span).to have_received(:set_tag).with(
+        Datadog::Tracing::Contrib::Aws::Ext::TAG_AWS_ACCOUNT,
+        '123456789012'
       )
     end
   end
@@ -28,12 +36,20 @@
     let(:state_machine_arn) { 'arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine' }
     let(:params) { { state_machine_arn: state_machine_arn } }
 
-    it 'sets the state_machine_name based on the state_machine_arn' do
+    it 'sets the state_machine_name and state_machine_arn based on the state_machine_arn' do
       step_functions.add_tags(span, params)
       expect(span).to have_received(:set_tag).with(
         Datadog::Tracing::Contrib::Aws::Ext::TAG_STATE_MACHINE_NAME,
         'MyStateMachine'
       )
+      expect(span).to have_received(:set_tag).with(
+        Datadog::Tracing::Contrib::Aws::Ext::TAG_STATE_MACHINE_ARN,
+        'arn:aws:states:us-east-1:123456789012:stateMachine:MyStateMachine'
+      )
+      expect(span).to have_received(:set_tag).with(
+        Datadog::Tracing::Contrib::Aws::Ext::TAG_AWS_ACCOUNT,
+        '123456789012'
+      )
     end
 
     it 'sets the state_machine_account_id based on the state_machine_arn' do