Skip to content

Commit

Permalink
Merge branch 'main' into patch-1
Browse files Browse the repository at this point in the history
  • Loading branch information
ocelotl authored Apr 22, 2024
2 parents 25d866f + c644f0d commit 2f45948
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 16 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,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

0 comments on commit 2f45948

Please sign in to comment.