botocore: Add support for AWS Step Functions semantic convention attributes

AWS Step Functions defines two semantic convention attributes:

aws.step_functions.activity.arn
aws.step_functions.state_machine.arn

https://github.com/open-telemetry/semantic-conventions/blob/main/docs/registry/attributes/aws.md#amazon-step-functions-attributes

Currently, these attributes are not set in the botocore instrumentation library. This PR adds support for them by extracting values from both Request and Response objects.

Tests
Added new unit tests (passing).

Verified with:
tox -e py312-test-instrumentation-botocore
tox -e spellcheck
tox -e lint-instrumentation-botocore
tox -e ruff

Backward Compatibility
This change is backward compatible. It only adds instrumentation for additional AWS resources and does not modify existing behavior in the auto-instrumentation library.

Note
This PR depends on #3736. Since #3736 has not yet been merged, its changes are included here as well.

Author: lukeina2z

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/__init__.py

@@ -35,6 +35,7 @@ def loader():
     "bedrock-runtime": _lazy_load(".bedrock", "_BedrockRuntimeExtension"),
     "dynamodb": _lazy_load(".dynamodb", "_DynamoDbExtension"),
     "lambda": _lazy_load(".lmbd", "_LambdaExtension"),
+    "stepfunctions": _lazy_load(".sfns", "_StepFunctionsExtension"),
     "sns": _lazy_load(".sns", "_SnsExtension"),
     "sqs": _lazy_load(".sqs", "_SqsExtension"),
 }

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/sfns.py

@@ -0,0 +1,58 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from opentelemetry.instrumentation.botocore.extensions.types import (
+    _AttributeMapT,
+    _AwsSdkExtension,
+    _BotocoreInstrumentorContext,
+    _BotoResultT,
+)
+from opentelemetry.semconv._incubating.attributes.aws_attributes import (
+    AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
+    AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
+)
+from opentelemetry.trace.span import Span
+
+
+class _StepFunctionsExtension(_AwsSdkExtension):
+    @staticmethod
+    def _set_arn_attributes(source, target, setter_func):
+        """Helper to set ARN attributes if they exist in source."""
+        activity_arn = source.get("activityArn")
+        if activity_arn:
+            setter_func(target, AWS_STEP_FUNCTIONS_ACTIVITY_ARN, activity_arn)
+
+        state_machine_arn = source.get("stateMachineArn")
+        if state_machine_arn:
+            setter_func(
+                target, AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN, state_machine_arn
+            )
+
+    def extract_attributes(self, attributes: _AttributeMapT):
+        self._set_arn_attributes(
+            self._call_context.params,
+            attributes,
+            lambda target, key, value: target.__setitem__(key, value),
+        )
+
+    def on_success(
+        self,
+        span: Span,
+        result: _BotoResultT,
+        instrumentor_context: _BotocoreInstrumentorContext,
+    ):
+        self._set_arn_attributes(
+            result,
+            span,
+            lambda target, key, value: target.set_attribute(key, value),
+        )

instrumentation/opentelemetry-instrumentation-botocore/test-requirements-0.txt

@@ -13,7 +13,7 @@ iniconfig==2.0.0
 Jinja2==3.1.6
 jmespath==1.0.1
 MarkupSafe==2.1.5
-moto==5.0.9
+moto==5.1.11
 packaging==24.0
 pluggy==1.5.0
 py-cpuinfo==9.0.0

instrumentation/opentelemetry-instrumentation-botocore/test-requirements-1.txt

@@ -13,7 +13,7 @@ iniconfig==2.0.0
 Jinja2==3.1.6
 jmespath==1.0.1
 MarkupSafe==2.1.5
-moto==5.0.9
+moto==5.1.11
 packaging==24.0
 pluggy==1.5.0
 py-cpuinfo==9.0.0

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_dynamodb.py

@@ -439,7 +439,7 @@ def test_scan(self):
             Limit=42,
             Select="ALL_ATTRIBUTES",
             TotalSegments=17,
-            Segment=21,
+            Segment=16,
             ProjectionExpression="PE",
             ConsistentRead=True,
             ReturnConsumedCapacity="TOTAL",
@@ -448,14 +448,14 @@ def test_scan(self):
         span = self.assert_span("Scan")
         self.assert_table_names(span, self.default_table_name)
         self.assertEqual(
-            21, span.attributes[SpanAttributes.AWS_DYNAMODB_SEGMENT]
+            16, span.attributes[SpanAttributes.AWS_DYNAMODB_SEGMENT]
         )
         self.assertEqual(
             17, span.attributes[SpanAttributes.AWS_DYNAMODB_TOTAL_SEGMENTS]
         )
-        self.assertEqual(1, span.attributes[SpanAttributes.AWS_DYNAMODB_COUNT])
+        self.assertEqual(0, span.attributes[SpanAttributes.AWS_DYNAMODB_COUNT])
         self.assertEqual(
-            1, span.attributes[SpanAttributes.AWS_DYNAMODB_SCANNED_COUNT]
+            0, span.attributes[SpanAttributes.AWS_DYNAMODB_SCANNED_COUNT]
         )
         self.assert_attributes_to_get(span, "id", "idl")
         self.assert_consistent_read(span, True)

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_sfns.py

@@ -0,0 +1,117 @@
+import json
+
+import botocore.session
+from moto import mock_aws
+
+from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
+from opentelemetry.semconv._incubating.attributes.aws_attributes import (
+    AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
+    AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
+)
+from opentelemetry.test.test_base import TestBase
+
+
+class TestSfnsExtension(TestBase):
+    def setUp(self):
+        super().setUp()
+        BotocoreInstrumentor().instrument()
+        session = botocore.session.get_session()
+        session.set_credentials(
+            access_key="access-key", secret_key="secret-key"
+        )
+        self.region = "us-west-2"
+        self.client = session.create_client(
+            "stepfunctions", region_name=self.region
+        )
+
+    def tearDown(self):
+        super().tearDown()
+        BotocoreInstrumentor().uninstrument()
+
+    SIMPLE_STATE_MACHINE_DEF = {
+        "Comment": "A simple Hello World example",
+        "StartAt": "HelloState",
+        "States": {
+            "HelloState": {
+                "Type": "Pass",
+                "Result": "Hello, Moto!",
+                "End": True,
+            }
+        },
+    }
+
+    def create_state_machine_and_get_arn(
+        self, name: str = "TestStateMachine"
+    ) -> str:
+        """
+        Create a state machine in mocked Step Functions and return its ARN.
+        """
+        definition_json = json.dumps(self.SIMPLE_STATE_MACHINE_DEF)
+        role_arn = "arn:aws:iam::123456789012:role/DummyRole"
+
+        response = self.client.create_state_machine(
+            name=name,
+            definition=definition_json,
+            roleArn=role_arn,
+        )
+        return response["stateMachineArn"]
+
+    def create_activity_and_get_arn(self, name: str = "TestActivity") -> str:
+        """
+        Create an activity in mocked Step Functions and return its ARN.
+        """
+        response = self.client.create_activity(
+            name=name,
+        )
+        return response["activityArn"]
+
+    @mock_aws
+    def test_sfns_create_state_machine(self):
+        state_machine_arn = self.create_state_machine_and_get_arn()
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(
+            span.attributes[AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN],
+            state_machine_arn,
+        )
+
+    @mock_aws
+    def test_sfns_describe_state_machine(self):
+        state_machine_arn = self.create_state_machine_and_get_arn()
+        self.client.describe_state_machine(stateMachineArn=state_machine_arn)
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        span = spans[1]
+        self.assertEqual(
+            span.attributes[AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN],
+            state_machine_arn,
+        )
+
+    @mock_aws
+    def test_sfns_create_activity(self):
+        activity_arn = self.create_activity_and_get_arn()
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(
+            span.attributes[AWS_STEP_FUNCTIONS_ACTIVITY_ARN],
+            activity_arn,
+        )
+
+    @mock_aws
+    def test_sfns_describe_activity(self):
+        activity_arn = self.create_activity_and_get_arn()
+        self.client.describe_activity(activityArn=activity_arn)
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        span = spans[1]
+        self.assertEqual(
+            span.attributes[AWS_STEP_FUNCTIONS_ACTIVITY_ARN],
+            activity_arn,
+        )