Skip to content

Commit

Permalink
[boto3sqs] Reduce number of wrapper targets
Browse files Browse the repository at this point in the history
There are actually 6 ways to initialize a boto3 API object.

```py
boto3.client()              # Using default global session
boto3.resource()            # Using default global session
boto3.Session().client()    # Using "re-exported" session.Session
boto3.Session().resource()  # Using "re-exported" session.Session

boto3.session.Session().client()    # Using session.Session directly
boto3.session.Session().resource()  # Using session.Session directly
```

We only have to patch `session.Session.client` to catch all the cases.

- https://github.com/boto/boto3/blob/b3c158c62aa2a1314dc0ec78caea1ea976abd1a0/boto3/session.py#L217-L229
- https://github.com/boto/boto3/blob/b3c158c62aa2a1314dc0ec78caea1ea976abd1a0/boto3/session.py#L446-L457
  • Loading branch information
mattoberle committed Feb 23, 2024
1 parent 0d7748b commit f634266
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 10 deletions.
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,10 +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, "resource", client_wrapper)
wrap_function_wrapper(boto3.Session, "client", client_wrapper)
wrap_function_wrapper(boto3.Session, "resource", client_wrapper)
wrap_function_wrapper(boto3.session.Session, "client", client_wrapper)

def _decorate_sqs(self, sqs_class: type) -> None:
"""
Expand Down Expand Up @@ -436,10 +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, "resource")
unwrap(boto3.Session, "client")
unwrap(boto3.Session, "resource")
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 @@ -57,7 +57,6 @@ def _make_sqs_resource(*, session=False):

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 @@ -66,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 @@ -80,12 +90,14 @@ def test_instrument_api_before_client_init(self) -> None:
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:
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):
for session in (False, True):
Expand All @@ -98,12 +110,14 @@ def test_instrument_api_before_resource_init(self) -> None:
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):
Expand Down

0 comments on commit f634266

Please sign in to comment.