a single hook for all aws services

Author: ItayGibel-helios

CHANGELOG.md

@@ -14,7 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#670](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/670))
 - `opentelemetry-instrumentation-redis` added request_hook and response_hook callbacks passed as arguments to the instrument method.
   ([#669](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/669))
-- `opentelemetry-instrumentation-botocore` add request_hooks and response_hooks
+- `opentelemetry-instrumentation-botocore` add `request_hook` and `response_hook` callbacks
   ([679](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/679))
 
 ### Changed

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

@@ -48,10 +48,10 @@
 The `instrument` method accepts the following keyword args:
 
 tracer_provider (TracerProvider) - an optional tracer provider
-request_hooks (dict) - a mapping between service names their respective callable request hooks
-* a request hook signature is: def request_hook(span: Span, operation_name: str, api_params: dict) -> None
-response_hooks (dict) - a mapping between service names their respective callable response hooks
-* a response hook signature is: def response_hook(span: Span, operation_name: str, result: dict) -> None
+request_hook (Callable) - a function with extra user-defined logic to be performed before performing the request
+this function signature is:  def request_hook(span: Span, service_name: str, operation_name: str, api_params: dict) -> None
+response_hook (Callable) - a function with extra user-defined logic to be performed after performing the request
+this function signature is:  def request_hook(span: Span, service_name: str, operation_name: str, result: dict) -> None
 
 for example:
 
@@ -60,16 +60,14 @@
     from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
     import botocore
 
-    def ec2_request_hook(span, operation_name, api_params):
+    def request_hook(span, service_name, operation_name, api_params):
         # request hook logic
 
-    def ec2_response_hook(span, operation_name, result):
+    def response_hook(span, service_name, operation_name, result):
         # response hook logic
 
     # Instrument Botocore with hooks
-    BotocoreInstrumentor().instrument(
-        request_hooks={"ec2": ec2_request_hook}, response_hooks={"ec2": ec2_response_hook}
-    )
+    BotocoreInstrumentor().instrument(request_hook=request_hook, response_hooks=response_hook)
 
     # This will create a span with Botocore-specific attributes, including custom attributes added from the hooks
     session = botocore.session.get_session()
@@ -127,8 +125,8 @@ class BotocoreInstrumentor(BaseInstrumentor):
 
     def __init__(self):
         super().__init__()
-        self.request_hooks = dict()
-        self.response_hooks = dict()
+        self.request_hook = None
+        self.response_hook = None
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -139,14 +137,8 @@ def _instrument(self, **kwargs):
             __name__, __version__, kwargs.get("tracer_provider")
         )
 
-        request_hooks = kwargs.get("request_hooks")
-        response_hooks = kwargs.get("response_hooks")
-
-        if isinstance(request_hooks, dict):
-            self.request_hooks = request_hooks
-
-        if isinstance(response_hooks, dict):
-            self.response_hooks = response_hooks
+        self.request_hook = kwargs.get("request_hook")
+        self.response_hook = kwargs.get("response_hook")
 
         wrap_function_wrapper(
             "botocore.client",
@@ -214,9 +206,10 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
                 context_api.set_value(_SUPPRESS_HTTP_INSTRUMENTATION_KEY, True)
             )
 
-            self.apply_request_hook(
-                span, service_name, operation_name, api_params
-            )
+            if callable(self.request_hook):
+                self.request_hook(
+                    span, service_name, operation_name, api_params
+                )
 
             try:
                 result = original_func(*args, **kwargs)
@@ -228,9 +221,8 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
             if error:
                 result = error.response
 
-            self.apply_response_hook(
-                span, service_name, operation_name, result
-            )
+            if callable(self.response_hook):
+                self.response_hook(span, service_name, operation_name, result)
 
             self._set_api_call_result_attributes(span, result)
 
@@ -284,17 +276,3 @@ def _set_api_call_result_attributes(span, result):
                         SpanAttributes.HTTP_STATUS_CODE,
                         metadata["HTTPStatusCode"],
                     )
-
-    def apply_request_hook(
-        self, span, service_name, operation_name, api_params
-    ):
-        if service_name in self.request_hooks:
-            request_hook = self.request_hooks.get(service_name)
-            if callable(request_hook):
-                request_hook(span, operation_name, api_params)
-
-    def apply_response_hook(self, span, service_name, operation_name, result):
-        if service_name in self.response_hooks:
-            response_hook = self.response_hooks.get(service_name)
-            if callable(response_hook):
-                response_hook(span, operation_name, result)

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -631,34 +631,22 @@ def test_dynamodb_client(self):
         )
 
     @mock_dynamodb2
-    def test_hooks(self):
+    def test__request_hook(self):
+        request_hook_service_attribute_name = "request_hook.service_name"
         request_hook_operation_attribute_name = "request_hook.operation_name"
         request_hook_api_params_attribute_name = "request_hook.api_params"
-        response_hook_operation_attribute_name = "response_hook.operation_name"
-        response_hook_result_attribute_name = "response_hook.result"
 
-        def request_hook(span, operation_name, api_params):
+        def request_hook(span, service_name, operation_name, api_params):
             hook_attributes = {
+                request_hook_service_attribute_name: service_name,
                 request_hook_operation_attribute_name: operation_name,
                 request_hook_api_params_attribute_name: json.dumps(api_params),
             }
             if span and span.is_recording():
                 span.set_attributes(hook_attributes)
 
-        def response_hook(span, operation_name, result):
-            if span and span.is_recording():
-                span.set_attribute(
-                    response_hook_operation_attribute_name, operation_name,
-                )
-                span.set_attribute(
-                    response_hook_result_attribute_name, list(result.keys()),
-                )
-
         BotocoreInstrumentor().uninstrument()
-        BotocoreInstrumentor().instrument(
-            request_hooks={"dynamodb": request_hook},
-            response_hooks={"dynamodb": response_hook},
-        )
+        BotocoreInstrumentor().instrument(request_hook=request_hook,)
 
         self.session = botocore.session.get_session()
         self.session.set_credentials(
@@ -694,8 +682,10 @@ def response_hook(span, operation_name, result):
             {"TableName": test_table_name, "Item": item}
         )
 
-        expected_result_keys = ("ConsumedCapacity", "ResponseMetadata")
-
+        self.assertEqual(
+            "dynamodb",
+            get_item_attributes.get(request_hook_service_attribute_name),
+        )
         self.assertEqual(
             "PutItem",
             get_item_attributes.get(request_hook_operation_attribute_name),
@@ -704,6 +694,61 @@ def response_hook(span, operation_name, result):
             expected_api_params,
             get_item_attributes.get(request_hook_api_params_attribute_name),
         )
+
+    @mock_dynamodb2
+    def test__response_hook(self):
+        response_hook_service_attribute_name = "request_hook.service_name"
+        response_hook_operation_attribute_name = "response_hook.operation_name"
+        response_hook_result_attribute_name = "response_hook.result"
+
+        def response_hook(span, service_name, operation_name, result):
+            hook_attributes = {
+                response_hook_service_attribute_name: service_name,
+                response_hook_operation_attribute_name: operation_name,
+                response_hook_result_attribute_name: list(result.keys()),
+            }
+            if span and span.is_recording():
+                span.set_attributes(hook_attributes)
+
+        BotocoreInstrumentor().uninstrument()
+        BotocoreInstrumentor().instrument(response_hook=response_hook,)
+
+        self.session = botocore.session.get_session()
+        self.session.set_credentials(
+            access_key="access-key", secret_key="secret-key"
+        )
+
+        ddb = self.session.create_client("dynamodb", region_name="us-west-2")
+
+        test_table_name = "test_table_name"
+
+        ddb.create_table(
+            AttributeDefinitions=[
+                {"AttributeName": "id", "AttributeType": "S"},
+            ],
+            KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
+            ProvisionedThroughput={
+                "ReadCapacityUnits": 5,
+                "WriteCapacityUnits": 5,
+            },
+            TableName=test_table_name,
+        )
+
+        item = {"id": {"S": "test_key"}}
+
+        ddb.put_item(TableName=test_table_name, Item=item)
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        get_item_attributes = spans[1].attributes
+
+        expected_result_keys = ("ConsumedCapacity", "ResponseMetadata")
+
+        self.assertEqual(
+            "dynamodb",
+            get_item_attributes.get(response_hook_service_attribute_name),
+        )
         self.assertEqual(
             "PutItem",
             get_item_attributes.get(response_hook_operation_attribute_name),