Skip to content

Commit 4afa1ee

Browse files
lukeina2zxrmx
andauthored
botocore: Add support for AWS Step Functions semantic convention attributes (#3737)
* 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. * add changelog. --------- Co-authored-by: Riccardo Magliocchetti <[email protected]>
1 parent 4db9772 commit 4afa1ee

File tree

4 files changed

+178
-0
lines changed

4 files changed

+178
-0
lines changed

CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
4848
([#3366](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3366))
4949
- `opentelemetry-instrumentation`: add support for `OTEL_PYTHON_AUTO_INSTRUMENTATION_EXPERIMENTAL_GEVENT_PATCH` to inform opentelemetry-instrument about gevent monkeypatching
5050
([#3699](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3699))
51+
- `opentelemetry-instrumentation`: botocore: Add support for AWS Step Functions semantic convention attributes
52+
([#3737](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3737))
5153
- `opentelemetry-instrumentation-botocore`: Add support for SNS semantic convention attribute aws.sns.topic.arn
5254
([#3734](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3734))
5355
- `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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ def loader():
3535
"bedrock-runtime": _lazy_load(".bedrock", "_BedrockRuntimeExtension"),
3636
"dynamodb": _lazy_load(".dynamodb", "_DynamoDbExtension"),
3737
"lambda": _lazy_load(".lmbd", "_LambdaExtension"),
38+
"stepfunctions": _lazy_load(".sfns", "_StepFunctionsExtension"),
3839
"sns": _lazy_load(".sns", "_SnsExtension"),
3940
"sqs": _lazy_load(".sqs", "_SqsExtension"),
4041
}
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
from opentelemetry.instrumentation.botocore.extensions.types import (
15+
_AttributeMapT,
16+
_AwsSdkExtension,
17+
_BotocoreInstrumentorContext,
18+
_BotoResultT,
19+
)
20+
from opentelemetry.semconv._incubating.attributes.aws_attributes import (
21+
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
22+
AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
23+
)
24+
from opentelemetry.trace.span import Span
25+
26+
27+
class _StepFunctionsExtension(_AwsSdkExtension):
28+
@staticmethod
29+
def _set_arn_attributes(source, target, setter_func):
30+
"""Helper to set ARN attributes if they exist in source."""
31+
activity_arn = source.get("activityArn")
32+
if activity_arn:
33+
setter_func(target, AWS_STEP_FUNCTIONS_ACTIVITY_ARN, activity_arn)
34+
35+
state_machine_arn = source.get("stateMachineArn")
36+
if state_machine_arn:
37+
setter_func(
38+
target, AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN, state_machine_arn
39+
)
40+
41+
def extract_attributes(self, attributes: _AttributeMapT):
42+
self._set_arn_attributes(
43+
self._call_context.params,
44+
attributes,
45+
lambda target, key, value: target.__setitem__(key, value),
46+
)
47+
48+
def on_success(
49+
self,
50+
span: Span,
51+
result: _BotoResultT,
52+
instrumentor_context: _BotocoreInstrumentorContext,
53+
):
54+
self._set_arn_attributes(
55+
result,
56+
span,
57+
lambda target, key, value: target.set_attribute(key, value),
58+
)
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
import json
2+
3+
import botocore.session
4+
from moto import mock_aws
5+
6+
from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
7+
from opentelemetry.semconv._incubating.attributes.aws_attributes import (
8+
AWS_STEP_FUNCTIONS_ACTIVITY_ARN,
9+
AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN,
10+
)
11+
from opentelemetry.test.test_base import TestBase
12+
13+
14+
class TestSfnsExtension(TestBase):
15+
def setUp(self):
16+
super().setUp()
17+
BotocoreInstrumentor().instrument()
18+
session = botocore.session.get_session()
19+
session.set_credentials(
20+
access_key="access-key", secret_key="secret-key"
21+
)
22+
self.region = "us-west-2"
23+
self.client = session.create_client(
24+
"stepfunctions", region_name=self.region
25+
)
26+
27+
def tearDown(self):
28+
super().tearDown()
29+
BotocoreInstrumentor().uninstrument()
30+
31+
SIMPLE_STATE_MACHINE_DEF = {
32+
"Comment": "A simple Hello World example",
33+
"StartAt": "HelloState",
34+
"States": {
35+
"HelloState": {
36+
"Type": "Pass",
37+
"Result": "Hello, Moto!",
38+
"End": True,
39+
}
40+
},
41+
}
42+
43+
def create_state_machine_and_get_arn(
44+
self, name: str = "TestStateMachine"
45+
) -> str:
46+
"""
47+
Create a state machine in mocked Step Functions and return its ARN.
48+
"""
49+
definition_json = json.dumps(self.SIMPLE_STATE_MACHINE_DEF)
50+
role_arn = "arn:aws:iam::123456789012:role/DummyRole"
51+
52+
response = self.client.create_state_machine(
53+
name=name,
54+
definition=definition_json,
55+
roleArn=role_arn,
56+
)
57+
return response["stateMachineArn"]
58+
59+
def create_activity_and_get_arn(self, name: str = "TestActivity") -> str:
60+
"""
61+
Create an activity in mocked Step Functions and return its ARN.
62+
"""
63+
response = self.client.create_activity(
64+
name=name,
65+
)
66+
return response["activityArn"]
67+
68+
@mock_aws
69+
def test_sfns_create_state_machine(self):
70+
state_machine_arn = self.create_state_machine_and_get_arn()
71+
spans = self.memory_exporter.get_finished_spans()
72+
assert spans
73+
self.assertEqual(len(spans), 1)
74+
span = spans[0]
75+
self.assertEqual(
76+
span.attributes[AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN],
77+
state_machine_arn,
78+
)
79+
80+
@mock_aws
81+
def test_sfns_describe_state_machine(self):
82+
state_machine_arn = self.create_state_machine_and_get_arn()
83+
self.client.describe_state_machine(stateMachineArn=state_machine_arn)
84+
85+
spans = self.memory_exporter.get_finished_spans()
86+
assert spans
87+
self.assertEqual(len(spans), 2)
88+
span = spans[1]
89+
self.assertEqual(
90+
span.attributes[AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN],
91+
state_machine_arn,
92+
)
93+
94+
@mock_aws
95+
def test_sfns_create_activity(self):
96+
activity_arn = self.create_activity_and_get_arn()
97+
spans = self.memory_exporter.get_finished_spans()
98+
assert spans
99+
self.assertEqual(len(spans), 1)
100+
span = spans[0]
101+
self.assertEqual(
102+
span.attributes[AWS_STEP_FUNCTIONS_ACTIVITY_ARN],
103+
activity_arn,
104+
)
105+
106+
@mock_aws
107+
def test_sfns_describe_activity(self):
108+
activity_arn = self.create_activity_and_get_arn()
109+
self.client.describe_activity(activityArn=activity_arn)
110+
spans = self.memory_exporter.get_finished_spans()
111+
assert spans
112+
self.assertEqual(len(spans), 2)
113+
span = spans[1]
114+
self.assertEqual(
115+
span.attributes[AWS_STEP_FUNCTIONS_ACTIVITY_ARN],
116+
activity_arn,
117+
)

0 commit comments

Comments
 (0)