Skip to content

Commit 2179fb9

Browse files
mariojonkemariojonke
authored
botocore: Add support for SNS publish and publish_batch (#1409)
1 parent 8dbd142 commit 2179fb9

File tree

6 files changed

+462
-0
lines changed

6 files changed

+462
-0
lines changed

‎CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
2424
([#1350](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1350))
2525
- `opentelemetry-instrumentation-starlette` Add support for regular expression matching and sanitization of HTTP headers.
2626
([#1404](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1404))
27+
- `opentelemetry-instrumentation-botocore` Add support for SNS `publish` and `publish_batch`.
28+
([#1409](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1409))
2729
- Strip leading comments from SQL queries when generating the span name.
2830
([#1434](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1434))
2931

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

+1
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ def loader():
3434
_KNOWN_EXTENSIONS = {
3535
"dynamodb": _lazy_load(".dynamodb", "_DynamoDbExtension"),
3636
"lambda": _lazy_load(".lmbd", "_LambdaExtension"),
37+
"sns": _lazy_load(".sns", "_SnsExtension"),
3738
"sqs": _lazy_load(".sqs", "_SqsExtension"),
3839
}
3940

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/_messaging.py

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
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+
15+
import logging
16+
from typing import Any, MutableMapping
17+
18+
from opentelemetry.propagate import get_global_textmap, inject
19+
from opentelemetry.propagators.textmap import CarrierT, Setter
20+
21+
_logger = logging.getLogger(__name__)
22+
23+
_MAX_MESSAGE_ATTRIBUTES = 10
24+
25+
26+
class MessageAttributesSetter(Setter[CarrierT]):
27+
def set(self, carrier: CarrierT, key: str, value: str):
28+
carrier[key] = {
29+
"DataType": "String",
30+
"StringValue": value,
31+
}
32+
33+
34+
message_attributes_setter = MessageAttributesSetter()
35+
36+
37+
def inject_propagation_context(
38+
carrier: MutableMapping[str, Any]
39+
) -> MutableMapping[str, Any]:
40+
if carrier is None:
41+
carrier = {}
42+
43+
fields = get_global_textmap().fields
44+
if len(carrier.keys()) + len(fields) <= _MAX_MESSAGE_ATTRIBUTES:
45+
inject(carrier, setter=message_attributes_setter)
46+
else:
47+
_logger.warning(
48+
"botocore instrumentation: cannot set context propagation on "
49+
"SQS/SNS message due to maximum amount of MessageAttributes"
50+
)
51+
52+
return carrier

‎instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/sns.py

+166
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
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+
15+
import abc
16+
import inspect
17+
from typing import Any, Dict, MutableMapping, Optional, Tuple
18+
19+
from opentelemetry.instrumentation.botocore.extensions._messaging import (
20+
inject_propagation_context,
21+
)
22+
from opentelemetry.instrumentation.botocore.extensions.types import (
23+
_AttributeMapT,
24+
_AwsSdkCallContext,
25+
_AwsSdkExtension,
26+
)
27+
from opentelemetry.semconv.trace import (
28+
MessagingDestinationKindValues,
29+
SpanAttributes,
30+
)
31+
from opentelemetry.trace import SpanKind
32+
from opentelemetry.trace.span import Span
33+
34+
################################################################################
35+
# SNS operations
36+
################################################################################
37+
38+
39+
class _SnsOperation(abc.ABC):
40+
@classmethod
41+
@abc.abstractmethod
42+
def operation_name(cls) -> str:
43+
pass
44+
45+
@classmethod
46+
def span_kind(cls) -> SpanKind:
47+
return SpanKind.CLIENT
48+
49+
@classmethod
50+
def extract_attributes(
51+
cls, call_context: _AwsSdkCallContext, attributes: _AttributeMapT
52+
):
53+
pass
54+
55+
@classmethod
56+
def before_service_call(cls, call_context: _AwsSdkCallContext, span: Span):
57+
pass
58+
59+
60+
class _OpPublish(_SnsOperation):
61+
_arn_arg_names = ("TopicArn", "TargetArn")
62+
_phone_arg_name = "PhoneNumber"
63+
64+
@classmethod
65+
def operation_name(cls) -> str:
66+
return "Publish"
67+
68+
@classmethod
69+
def span_kind(cls) -> SpanKind:
70+
return SpanKind.PRODUCER
71+
72+
@classmethod
73+
def extract_attributes(
74+
cls, call_context: _AwsSdkCallContext, attributes: _AttributeMapT
75+
):
76+
destination_name, is_phone_number = cls._extract_destination_name(
77+
call_context
78+
)
79+
attributes[
80+
SpanAttributes.MESSAGING_DESTINATION_KIND
81+
] = MessagingDestinationKindValues.TOPIC.value
82+
attributes[SpanAttributes.MESSAGING_DESTINATION] = destination_name
83+
84+
call_context.span_name = (
85+
f"{'phone_number' if is_phone_number else destination_name} send"
86+
)
87+
88+
@classmethod
89+
def _extract_destination_name(
90+
cls, call_context: _AwsSdkCallContext
91+
) -> Tuple[str, bool]:
92+
arn = cls._extract_input_arn(call_context)
93+
if arn:
94+
return arn.rsplit(":", 1)[-1], False
95+
96+
if cls._phone_arg_name:
97+
phone_number = call_context.params.get(cls._phone_arg_name)
98+
if phone_number:
99+
return phone_number, True
100+
101+
return "unknown", False
102+
103+
@classmethod
104+
def _extract_input_arn(
105+
cls, call_context: _AwsSdkCallContext
106+
) -> Optional[str]:
107+
for input_arn in cls._arn_arg_names:
108+
arn = call_context.params.get(input_arn)
109+
if arn:
110+
return arn
111+
return None
112+
113+
@classmethod
114+
def before_service_call(cls, call_context: _AwsSdkCallContext, span: Span):
115+
cls._inject_span_into_entry(call_context.params)
116+
117+
@classmethod
118+
def _inject_span_into_entry(cls, entry: MutableMapping[str, Any]):
119+
entry["MessageAttributes"] = inject_propagation_context(
120+
entry.get("MessageAttributes")
121+
)
122+
123+
124+
class _OpPublishBatch(_OpPublish):
125+
_arn_arg_names = ("TopicArn",)
126+
_phone_arg_name = None
127+
128+
@classmethod
129+
def operation_name(cls) -> str:
130+
return "PublishBatch"
131+
132+
@classmethod
133+
def before_service_call(cls, call_context: _AwsSdkCallContext, span: Span):
134+
for entry in call_context.params.get("PublishBatchRequestEntries", ()):
135+
cls._inject_span_into_entry(entry)
136+
137+
138+
################################################################################
139+
# SNS extension
140+
################################################################################
141+
142+
_OPERATION_MAPPING = {
143+
op.operation_name(): op
144+
for op in globals().values()
145+
if inspect.isclass(op)
146+
and issubclass(op, _SnsOperation)
147+
and not inspect.isabstract(op)
148+
} # type: Dict[str, _SnsOperation]
149+
150+
151+
class _SnsExtension(_AwsSdkExtension):
152+
def __init__(self, call_context: _AwsSdkCallContext):
153+
super().__init__(call_context)
154+
self._op = _OPERATION_MAPPING.get(call_context.operation)
155+
if self._op:
156+
call_context.span_kind = self._op.span_kind()
157+
158+
def extract_attributes(self, attributes: _AttributeMapT):
159+
attributes[SpanAttributes.MESSAGING_SYSTEM] = "aws.sns"
160+
161+
if self._op:
162+
self._op.extract_attributes(self._call_context, attributes)
163+
164+
def before_service_call(self, span: Span):
165+
if self._op:
166+
self._op.before_service_call(self._call_context, span)

‎instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_messaging.py

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
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+
15+
from opentelemetry.instrumentation.botocore.extensions._messaging import (
16+
inject_propagation_context,
17+
message_attributes_setter,
18+
)
19+
from opentelemetry.test.test_base import TestBase
20+
21+
22+
class TestMessageAttributes(TestBase):
23+
def test_message_attributes_setter(self):
24+
carrier = {}
25+
26+
message_attributes_setter.set(carrier, "key", "value")
27+
self.assertEqual(
28+
{"key": {"DataType": "String", "StringValue": "value"}}, carrier
29+
)
30+
31+
def test_inject_propagation_context(self):
32+
carrier = {
33+
"key1": {"DataType": "String", "StringValue": "value1"},
34+
"key2": {"DataType": "String", "StringValue": "value2"},
35+
}
36+
37+
tracer = self.tracer_provider.get_tracer("test-tracer")
38+
with tracer.start_as_current_span("span"):
39+
inject_propagation_context(carrier)
40+
41+
self.assertGreater(len(carrier), 2)
42+
43+
def test_inject_propagation_context_too_many_attributes(self):
44+
carrier = {
45+
f"key{idx}": {"DataType": "String", "StringValue": f"value{idx}"}
46+
for idx in range(10)
47+
}
48+
tracer = self.tracer_provider.get_tracer("test-tracer")
49+
with tracer.start_as_current_span("span"):
50+
inject_propagation_context(carrier)
51+
52+
self.assertEqual(10, len(carrier))

0 commit comments

Comments
 (0)