botocore: Add support for AWS Step Functions semantic convention attributes

CHANGELOG.md

@@ -48,6 +48,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#3366](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3366))
 - `opentelemetry-instrumentation`: add support for  `OTEL_PYTHON_AUTO_INSTRUMENTATION_EXPERIMENTAL_GEVENT_PATCH` to inform opentelemetry-instrument about gevent monkeypatching
   ([#3699](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3699))
+- `opentelemetry-instrumentation`: botocore: Add support for AWS Step Functions semantic convention attributes
+  ([#3737](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3737))
 - `opentelemetry-instrumentation-botocore`: Add support for SNS semantic convention attribute aws.sns.topic.arn
   ([#3734](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3734))
 - `opentelemetry-instrumentation`: botocore: upgrade moto package from 5.0.9 to 5.1.11

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/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,
+        )