feat: Add auto-instrumentation support for StepFunctions

Author: yiyuan-he

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -25,3 +25,5 @@
 AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
 AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn"
 AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
+AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
+AWS_STEPFUNCTIONS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -24,6 +24,8 @@
     AWS_SPAN_KIND,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
     GEN_AI_REQUEST_MODEL,
@@ -92,6 +94,7 @@
 _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime"
 _NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager"
 _NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS"
+_NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -315,6 +318,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str
             "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME,
             "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME,
             "SNS": _NORMALIZED_SNS_SERVICE_NAME,
+            "SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME,
         }
         return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name)
     return service_name
@@ -428,6 +432,12 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         elif is_key_present(span, AWS_SNS_TOPIC_ARN):
             remote_resource_type = _NORMALIZED_SNS_SERVICE_NAME + "::Topic"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SNS_TOPIC_ARN))
+        elif is_key_present(span, AWS_STEPFUNCTIONS_STATEMACHINE_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::StateMachine"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_STATEMACHINE_ARN))
+        elif is_key_present(span, AWS_STEPFUNCTIONS_ACTIVITY_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::Activity"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -9,6 +9,8 @@
     AWS_SNS_TOPIC_ARN,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro.patches._bedrock_patches import (  # noqa # pylint: disable=unused-import
     _BedrockAgentExtension,
@@ -35,6 +37,20 @@ def _apply_botocore_instrumentation_patches() -> None:
     _apply_botocore_bedrock_patch()
     _apply_botocore_secretsmanager_patch()
     _apply_botocore_sns_patch()
+    _apply_botocore_stepfunctions_patch()
+
+
+def _apply_botocore_stepfunctions_patch() -> None:
+    """Botocore instrumentation patch for StepFunctions
+
+    This patch adds an extension to the upstream's list of known extensions for 
+    StepFunctions. Extensions allow for custom logic for adding service-specific
+    information to spans, such as attributes. Specifically, we are adding logic
+    to add the `aws.stepfunctions.state_machine.arn` and `aws.stepfunctions.activity.arn`
+    attributes, to be used to generate RemoteTarget and achieve partity with the
+    Java instrumentation.
+    """
+    _KNOWN_EXTENSIONS["stepfunctions"] = _lazy_load(".", "_StepFunctionsExtension")
 
 
 def _apply_botocore_sns_patch() -> None:
@@ -149,6 +165,16 @@ def loader():
 # END The OpenTelemetry Authors code
 
 
+class _StepFunctionsExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        state_machine_arn = self._call_context.params.get("stateMachineArn")
+        if state_machine_arn:
+            attributes[AWS_STEPFUNCTIONS_STATEMACHINE_ARN] = state_machine_arn
+        activity_arn = self._call_context.params.get("activityArn")
+        if activity_arn:
+            attributes[AWS_STEPFUNCTIONS_ACTIVITY_ARN] = activity_arn
+
+
 class _SecretsManagerExtension(_AwsSdkExtension):
     def extract_attributes(self, attributes: _AttributeMapT):
         """

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py

@@ -26,6 +26,8 @@
     AWS_SPAN_KIND,
     AWS_SQS_QUEUE_NAME,
     AWS_SQS_QUEUE_URL,
+    AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+    AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
 )
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
 from amazon.opentelemetry.distro._aws_span_processing_util import GEN_AI_REQUEST_MODEL
@@ -881,6 +883,7 @@ def test_normalize_remote_service_name_aws_sdk(self):
         self.validate_aws_sdk_service_normalization("Bedrock Runtime", "AWS::BedrockRuntime")
         self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager")
         self.validate_aws_sdk_service_normalization("SNS", "AWS::SNS")
+        self.validate_aws_sdk_service_normalization("SFN", "AWS::StepFunctions")
 
     def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str):
         self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name])
@@ -1114,6 +1117,31 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::SNS::Topic", "arn:aws:sns:us-west-2:012345678901:test_topic")
         self._mock_attribute([AWS_SNS_TOPIC_ARN], [None])
 
+        # Validate behaviour of AWS_STEPFUNCTIONS_STATEMACHINE_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_STEPFUNCTIONS_STATEMACHINE_ARN], 
+            ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"], 
+             keys, 
+             values
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::StepFunctions::StateMachine",
+            "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"
+        )
+        self._mock_attribute([AWS_STEPFUNCTIONS_STATEMACHINE_ARN], [None])
+
+        # Validate behaviour of AWS_STEPFUNCTIONS_ACTIVITY_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_STEPFUNCTIONS_ACTIVITY_ARN],
+            ["arn:aws:states:us-east-1:007003123456789012:activity:testActivity"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::StepFunctions::Activity", "arn:aws:states:us-east-1:007003123456789012:activity:testActivity"
+        )
+        self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py

@@ -28,6 +28,8 @@
 _GEN_AI_REQUEST_MODEL: str = "genAiReuqestModelId"
 _SECRET_ARN: str = "arn:aws:secretsmanager:us-west-2:000000000000:secret:testSecret-ABCDEF"
 _TOPIC_ARN: str = "topicArn"
+_STATE_MACHINE_ARN: str = "arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine"
+_ACTIVITY_ARN: str = "arn:aws:states:us-east-1:007003123456789012:activity:testActivity"
 
 # Patch names
 GET_DISTRIBUTION_PATCH: str = (
@@ -149,6 +151,9 @@ def _test_unpatched_botocore_instrumentation(self):
         # SNS
         self.assertTrue("sns" in _KNOWN_EXTENSIONS, "Upstream has removed the SNS extension")
 
+        # StepFunctions
+        self.assertFalse("stepfunctions" in _KNOWN_EXTENSIONS, "Upstream has added a StepFunctions extension")
+
     def _test_unpatched_gevent_instrumentation(self):
         self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched")
         self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched")
@@ -223,6 +228,14 @@ def _test_patched_botocore_instrumentation(self):
         self.assertTrue("aws.sns.topic.arn" in sns_attributes)
         self.assertEqual(sns_attributes["aws.sns.topic.arn"], _TOPIC_ARN)
 
+        # StepFunctions
+        self.assertTrue("stepfunctions" in _KNOWN_EXTENSIONS)
+        stepfunctions_attributes: Dict[str, str] = _do_extract_stepfunctions_attributes()
+        self.assertTrue("aws.stepfunctions.state_machine.arn" in stepfunctions_attributes)
+        self.assertEqual(stepfunctions_attributes["aws.stepfunctions.state_machine.arn"], _STATE_MACHINE_ARN)
+        self.assertTrue("aws.stepfunctions.activity.arn" in stepfunctions_attributes)
+        self.assertEqual(stepfunctions_attributes["aws.stepfunctions.activity.arn"], _ACTIVITY_ARN)
+
     def _test_patched_gevent_os_ssl_instrumentation(self):
         # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl'
         self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched")
@@ -399,6 +412,12 @@ def _do_extract_sns_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
+def _do_extract_stepfunctions_attributes() -> Dict[str, str]:
+    service_name: str = "stepfunctions"
+    params: Dict[str, str] = {"stateMachineArn": _STATE_MACHINE_ARN, "activityArn": _ACTIVITY_ARN}
+    return _do_extract_attributes(service_name, params)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params