Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[boto3sqs] Instrument Session and resource #2161

Merged
merged 11 commits into from
Apr 22, 2024
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

- `opentelemetry-instrumentation-grpc` AioClientInterceptor should propagate with a Metadata object
([#2363](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2363))
- `opentelemetry-instrumentation-boto3sqs` Instrument Session and resource
([#2161](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2161))

### Added

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
import logging
from typing import Any, Collection, Dict, Generator, List, Mapping, Optional

import boto3
import boto3.session
import botocore.client
from wrapt import wrap_function_wrapper

Expand Down Expand Up @@ -382,7 +382,7 @@ def client_wrapper(wrapped, instance, args, kwargs):
self._decorate_sqs(type(retval))
return retval

wrap_function_wrapper(boto3, "client", client_wrapper)
wrap_function_wrapper(boto3.session.Session, "client", client_wrapper)

def _decorate_sqs(self, sqs_class: type) -> None:
"""
Expand Down Expand Up @@ -433,7 +433,7 @@ def _instrument(self, **kwargs: Dict[str, Any]) -> None:
self._decorate_sqs(client_cls)

def _uninstrument(self, **kwargs: Dict[str, Any]) -> None:
unwrap(boto3, "client")
unwrap(boto3.session.Session, "client")

for client_cls in botocore.client.BaseClient.__subclasses__():
self._un_decorate_sqs(client_cls)
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

import boto3
from botocore.awsrequest import AWSResponse
from wrapt import BoundFunctionWrapper, FunctionWrapper
from wrapt import BoundFunctionWrapper

from opentelemetry.instrumentation.boto3sqs import (
Boto3SQSGetter,
Expand All @@ -37,8 +37,17 @@
from opentelemetry.trace.span import Span, format_span_id, format_trace_id


def _make_sqs_client():
return boto3.client(
def _make_sqs_client(*, session=False):
return (boto3.Session() if session else boto3).client(
"sqs",
region_name="us-east-1",
aws_access_key_id="dummy",
aws_secret_access_key="dummy",
)


def _make_sqs_resource(*, session=False):
return (boto3.Session() if session else boto3).resource(
"sqs",
region_name="us-east-1",
aws_access_key_id="dummy",
Expand All @@ -48,7 +57,6 @@ def _make_sqs_client():

class TestBoto3SQSInstrumentor(TestCase):
def _assert_instrumented(self, client):
self.assertIsInstance(boto3.client, FunctionWrapper)
self.assertIsInstance(client.send_message, BoundFunctionWrapper)
self.assertIsInstance(client.send_message_batch, BoundFunctionWrapper)
self.assertIsInstance(client.receive_message, BoundFunctionWrapper)
Expand All @@ -57,6 +65,17 @@ def _assert_instrumented(self, client):
client.delete_message_batch, BoundFunctionWrapper
)

def _assert_uninstrumented(self, client):
self.assertNotIsInstance(client.send_message, BoundFunctionWrapper)
self.assertNotIsInstance(
client.send_message_batch, BoundFunctionWrapper
)
self.assertNotIsInstance(client.receive_message, BoundFunctionWrapper)
self.assertNotIsInstance(client.delete_message, BoundFunctionWrapper)
self.assertNotIsInstance(
client.delete_message_batch, BoundFunctionWrapper
)

@staticmethod
@contextmanager
def _active_instrumentor():
Expand All @@ -67,19 +86,48 @@ def _active_instrumentor():
Boto3SQSInstrumentor().uninstrument()

def test_instrument_api_before_client_init(self) -> None:
with self._active_instrumentor():
client = _make_sqs_client()
self._assert_instrumented(client)
for session in (False, True):
with self._active_instrumentor():
client = _make_sqs_client(session=session)
self._assert_instrumented(client)
self._assert_uninstrumented(client)

def test_instrument_api_after_client_init(self) -> None:
client = _make_sqs_client()
with self._active_instrumentor():
self._assert_instrumented(client)
for session in (False, True):
client = _make_sqs_client(session=session)
with self._active_instrumentor():
self._assert_instrumented(client)
self._assert_uninstrumented(client)

def test_instrument_multiple_clients(self):
with self._active_instrumentor():
self._assert_instrumented(_make_sqs_client())
self._assert_instrumented(_make_sqs_client())
for session in (False, True):
with self._active_instrumentor():
self._assert_instrumented(_make_sqs_client(session=session))
self._assert_instrumented(_make_sqs_client(session=session))

def test_instrument_api_before_resource_init(self) -> None:
for session in (False, True):
with self._active_instrumentor():
sqs = _make_sqs_resource(session=session)
self._assert_instrumented(sqs.meta.client)
self._assert_uninstrumented(sqs.meta.client)

def test_instrument_api_after_resource_init(self) -> None:
for session in (False, True):
sqs = _make_sqs_resource(session=session)
with self._active_instrumentor():
self._assert_instrumented(sqs.meta.client)
self._assert_uninstrumented(sqs.meta.client)

def test_instrument_multiple_resources(self):
for session in (False, True):
with self._active_instrumentor():
self._assert_instrumented(
_make_sqs_resource(session=session).meta.client
)
self._assert_instrumented(
_make_sqs_resource(session=session).meta.client
)


class TestBoto3SQSGetter(TestCase):
Expand Down
Loading