feat: option to disallow tracer to capture response as metadata

CHANGELOG.md

@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
     - This is a breaking change if you were graphing/alerting on both application metrics with the same name to compensate this previous malfunctioning
     - Marked as bugfix as this is the intended behaviour since the beginning, as you shouldn't have the same application metric with different dimensions
 
+### Added
+- **Tracer**: capture_lambda_handler and capture_method decorators now support `capture_response` parameter to not include function's response as part of tracing metadata
+
 ## [1.3.1] - 2020-08-22
 ### Fixed
 - **Tracer**: capture_method decorator did not properly handle nested context managers

aws_lambda_powertools/tracing/tracer.py

@@ -150,27 +150,27 @@ def __init__(
         self.auto_patch = self._config["auto_patch"]
 
         if self.disabled:
-            self.__disable_tracing_provider()
+            self._disable_tracer_provider()
 
         if self.auto_patch:
             self.patch(modules=patch_modules)
 
     def put_annotation(self, key: str, value: Any):
         """Adds annotation to existing segment or subsegment
 
+        Parameters
+        ----------
+        key : str
+            Annotation key
+        value : any
+            Value for annotation
+
         Example
         -------
         Custom annotation for a pseudo service named payment
 
             tracer = Tracer(service="payment")
             tracer.put_annotation("PaymentStatus", "CONFIRMED")
-
-        Parameters
-        ----------
-        key : str
-            Annotation key (e.g. PaymentStatus)
-        value : any
-            Value for annotation (e.g. "CONFIRMED")
         """
         if self.disabled:
             logger.debug("Tracing has been disabled, aborting put_annotation")
@@ -226,12 +226,19 @@ def patch(self, modules: Tuple[str] = None):
         else:
             aws_xray_sdk.core.patch(modules)
 
-    def capture_lambda_handler(self, lambda_handler: Callable[[Dict, Any], Any] = None):
+    def capture_lambda_handler(self, lambda_handler: Callable[[Dict, Any], Any] = None, capture_response: bool = True):
         """Decorator to create subsegment for lambda handlers
 
         As Lambda follows (event, context) signature we can remove some of the boilerplate
         and also capture any exception any Lambda function throws or its response as metadata
 
+        Parameters
+        ----------
+        lambda_handler : Callable
+            Method to annotate on
+        capture_response : bool, optional
+            Instructs tracer to not include handler's response as metadata, by default True
+
         Example
         -------
         **Lambda function using capture_lambda_handler decorator**
@@ -241,16 +248,24 @@ def capture_lambda_handler(self, lambda_handler: Callable[[Dict, Any], Any] = No
             def handler(event, context):
                 ...
 
-        Parameters
-        ----------
-        method : Callable
-            Method to annotate on
+        **Preventing Tracer to log response as metadata**
+
+            tracer = Tracer(service="payment")
+            @tracer.capture_lambda_handler(capture_response=False)
+            def handler(event, context):
+                ...
 
         Raises
         ------
         err
             Exception raised by method
         """
+        # If handler is None we've been called with parameters
+        # Return a partial function with args filled
+        if lambda_handler is None:
+            logger.debug("Decorator called with parameters")
+            return functools.partial(self.capture_lambda_handler, capture_response=capture_response)
+
         lambda_handler_name = lambda_handler.__name__
 
         @functools.wraps(lambda_handler)
@@ -266,22 +281,24 @@ def decorate(event, context):
                     logger.debug("Calling lambda handler")
                     response = lambda_handler(event, context)
                     logger.debug("Received lambda handler response successfully")
-                    logger.debug(response)
                     self._add_response_as_metadata(
-                        function_name=lambda_handler_name, data=response, subsegment=subsegment
+                        method_name=lambda_handler_name,
+                        data=response,
+                        subsegment=subsegment,
+                        capture_response=capture_response,
                     )
                 except Exception as err:
                     logger.exception(f"Exception received from {lambda_handler_name}")
                     self._add_full_exception_as_metadata(
-                        function_name=lambda_handler_name, error=err, subsegment=subsegment
+                        method_name=lambda_handler_name, error=err, subsegment=subsegment
                     )
                     raise
 
                 return response
 
         return decorate
 
-    def capture_method(self, method: Callable = None):
+    def capture_method(self, method: Callable = None, capture_response: bool = True):
         """Decorator to create subsegment for arbitrary functions
 
         It also captures both response and exceptions as metadata
@@ -295,6 +312,13 @@ def capture_method(self, method: Callable = None):
         `async.gather` is called, or use `in_subsegment_async`
         context manager via our escape hatch mechanism - See examples.
 
+        Parameters
+        ----------
+        method : Callable
+            Method to annotate on
+        capture_response : bool, optional
+            Instructs tracer to not include method's response as metadata, by default True
+
         Example
         -------
         **Custom function using capture_method decorator**
@@ -416,69 +440,84 @@ async def async_tasks():
 
                 return { "task": "done", **ret }
 
-        Parameters
-        ----------
-        method : Callable
-            Method to annotate on
-
         Raises
         ------
         err
             Exception raised by method
         """
+        # If method is None we've been called with parameters
+        # Return a partial function with args filled
+        if method is None:
+            logger.debug("Decorator called with parameters")
+            return functools.partial(self.capture_method, capture_response=capture_response)
+
+        method_name = f"{method.__name__}"
 
         if inspect.iscoroutinefunction(method):
-            decorate = self._decorate_async_function(method=method)
+            decorate = self._decorate_async_function(
+                method=method, capture_response=capture_response, method_name=method_name
+            )
         elif inspect.isgeneratorfunction(method):
-            decorate = self._decorate_generator_function(method=method)
+            decorate = self._decorate_generator_function(
+                method=method, capture_response=capture_response, method_name=method_name
+            )
         elif hasattr(method, "__wrapped__") and inspect.isgeneratorfunction(method.__wrapped__):
-            decorate = self._decorate_generator_function_with_context_manager(method=method)
+            decorate = self._decorate_generator_function_with_context_manager(
+                method=method, capture_response=capture_response, method_name=method_name
+            )
         else:
-            decorate = self._decorate_sync_function(method=method)
+            decorate = self._decorate_sync_function(
+                method=method, capture_response=capture_response, method_name=method_name
+            )
 
         return decorate
 
-    def _decorate_async_function(self, method: Callable = None):
-        method_name = f"{method.__name__}"
-
+    def _decorate_async_function(self, method: Callable = None, capture_response: bool = True, method_name: str = None):
         @functools.wraps(method)
         async def decorate(*args, **kwargs):
             async with self.provider.in_subsegment_async(name=f"## {method_name}") as subsegment:
                 try:
                     logger.debug(f"Calling method: {method_name}")
                     response = await method(*args, **kwargs)
-                    self._add_response_as_metadata(function_name=method_name, data=response, subsegment=subsegment)
+                    self._add_response_as_metadata(
+                        method_name=method_name,
+                        data=response,
+                        subsegment=subsegment,
+                        capture_response=capture_response,
+                    )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
                     raise
 
                 return response
 
         return decorate
 
-    def _decorate_generator_function(self, method: Callable = None):
-        method_name = f"{method.__name__}"
-
+    def _decorate_generator_function(
+        self, method: Callable = None, capture_response: bool = True, method_name: str = None
+    ):
         @functools.wraps(method)
         def decorate(*args, **kwargs):
             with self.provider.in_subsegment(name=f"## {method_name}") as subsegment:
                 try:
                     logger.debug(f"Calling method: {method_name}")
                     result = yield from method(*args, **kwargs)
-                    self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
+                    self._add_response_as_metadata(
+                        method_name=method_name, data=result, subsegment=subsegment, capture_response=capture_response
+                    )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
                     raise
 
                 return result
 
         return decorate
 
-    def _decorate_generator_function_with_context_manager(self, method: Callable = None):
-        method_name = f"{method.__name__}"
-
+    def _decorate_generator_function_with_context_manager(
+        self, method: Callable = None, capture_response: bool = True, method_name: str = None
+    ):
         @functools.wraps(method)
         @contextlib.contextmanager
         def decorate(*args, **kwargs):
@@ -488,74 +527,87 @@ def decorate(*args, **kwargs):
                     with method(*args, **kwargs) as return_val:
                         result = return_val
                         yield result
-                    self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
+                    self._add_response_as_metadata(
+                        method_name=method_name, data=result, subsegment=subsegment, capture_response=capture_response
+                    )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
                     raise
 
         return decorate
 
-    def _decorate_sync_function(self, method: Callable = None):
-        method_name = f"{method.__name__}"
-
+    def _decorate_sync_function(self, method: Callable = None, capture_response: bool = True, method_name: str = None):
         @functools.wraps(method)
         def decorate(*args, **kwargs):
             with self.provider.in_subsegment(name=f"## {method_name}") as subsegment:
                 try:
                     logger.debug(f"Calling method: {method_name}")
                     response = method(*args, **kwargs)
-                    self._add_response_as_metadata(function_name=method_name, data=response, subsegment=subsegment)
+                    self._add_response_as_metadata(
+                        method_name=method_name,
+                        data=response,
+                        subsegment=subsegment,
+                        capture_response=capture_response,
+                    )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
                     raise
 
                 return response
 
         return decorate
 
     def _add_response_as_metadata(
-        self, function_name: str = None, data: Any = None, subsegment: aws_xray_sdk.core.models.subsegment = None
+        self,
+        method_name: str = None,
+        data: Any = None,
+        subsegment: aws_xray_sdk.core.models.subsegment = None,
+        capture_response: bool = True,
     ):
         """Add response as metadata for given subsegment
 
         Parameters
         ----------
-        function_name : str, optional
-            function name to add as metadata key, by default None
+        method_name : str, optional
+            method name to add as metadata key, by default None
         data : Any, optional
             data to add as subsegment metadata, by default None
         subsegment : aws_xray_sdk.core.models.subsegment, optional
             existing subsegment to add metadata on, by default None
+        capture_response : bool, optional
+            Do not include response as metadata, by default True
         """
-        if data is None or subsegment is None:
+        if data is None or not capture_response or subsegment is None:
             return
 
-        subsegment.put_metadata(key=f"{function_name} response", value=data, namespace=self._config["service"])
+        subsegment.put_metadata(key=f"{method_name} response", value=data, namespace=self._config["service"])
 
     def _add_full_exception_as_metadata(
-        self, function_name: str = None, error: Exception = None, subsegment: aws_xray_sdk.core.models.subsegment = None
+        self, method_name: str = None, error: Exception = None, subsegment: aws_xray_sdk.core.models.subsegment = None
     ):
         """Add full exception object as metadata for given subsegment
 
         Parameters
         ----------
-        function_name : str, optional
-            function name to add as metadata key, by default None
+        method_name : str, optional
+            method name to add as metadata key, by default None
         error : Exception, optional
             error to add as subsegment metadata, by default None
         subsegment : aws_xray_sdk.core.models.subsegment, optional
             existing subsegment to add metadata on, by default None
         """
-        subsegment.put_metadata(key=f"{function_name} error", value=error, namespace=self._config["service"])
+        subsegment.put_metadata(key=f"{method_name} error", value=error, namespace=self._config["service"])
 
-    def __disable_tracing_provider(self):
+    @staticmethod
+    def _disable_tracer_provider():
         """Forcefully disables tracing"""
         logger.debug("Disabling tracer provider...")
         aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
 
-    def __is_trace_disabled(self) -> bool:
+    @staticmethod
+    def _is_tracer_disabled() -> bool:
         """Detects whether trace has been disabled
 
         Tracing is automatically disabled in the following conditions:
@@ -592,7 +644,7 @@ def __build_config(
         provider: aws_xray_sdk.core.xray_recorder = None,
     ):
         """ Populates Tracer config for new and existing initializations """
-        is_disabled = disabled if disabled is not None else self.__is_trace_disabled()
+        is_disabled = disabled if disabled is not None else self._is_tracer_disabled()
         is_service = service if service is not None else os.getenv("POWERTOOLS_SERVICE_NAME")
 
         self._config["provider"] = provider if provider is not None else self._config["provider"]

docs/content/core/tracer.mdx

@@ -17,6 +17,12 @@ Tracer is an opinionated thin wrapper for [AWS X-Ray Python SDK](https://github.
 * Support tracing async methods, generators, and context managers
 * Auto patch supported modules, or a tuple of explicit modules supported by AWS X-Ray
 
+<Note type="warning">
+  <strong>Returning sensitive information from your Lambda handler or functions, where Tracer is used?</strong>
+  <br/><br/>
+  You can disable Tracer from capturing their responses as tracing metadata with <strong><code>capture_response=False</code></strong> parameter in both capture_lambda_handler and capture_method decorators.
+</Note><br/>
+
 ## Initialization
 
 Your AWS Lambda function must have permission to send traces to AWS X-Ray - Here is an example using AWS Serverless Application Model (SAM)
@@ -63,6 +69,10 @@ def handler(event, context):
     charge_id = event.get('charge_id')
     payment = collect_payment(charge_id)
     ...
+
+@tracer.capture_lambda_handler(capture_response=False) # highlight-line
+def handler(event, context):
+	return "sensitive_information"
 ```
 
 ### Annotations
@@ -108,7 +118,10 @@ def collect_payment(charge_id):
     ret = requests.post(PAYMENT_ENDPOINT) # logic
     tracer.put_annotation("PAYMENT_STATUS", "SUCCESS") # custom annotation
     return ret
-...
+
+@tracer.capture_method(capture_response=False) # highlight-line
+def sensitive_information_to_be_processed():
+	return "sensitive_information"
 ```
 
 ## Asynchronous and generator functions

tests/unit/test_tracing.py

@@ -179,10 +179,9 @@ def test_tracer_no_autopatch(patch_mock):
     assert patch_mock.call_count == 0
 
 
-def test_tracer_lambda_handler_does_not_add_empty_response_as_metadata(mocker, provider_stub):
+def test_tracer_lambda_handler_does_not_add_empty_response_as_metadata(mocker, provider_stub, in_subsegment_mock):
     # GIVEN tracer is initialized
-    put_metadata_mock = mocker.MagicMock()
-    provider = provider_stub(put_metadata_mock=put_metadata_mock)
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
     tracer = Tracer(provider=provider)
 
     # WHEN capture_lambda_handler decorator is used
@@ -194,13 +193,12 @@ def handler(event, context):
     handler({}, mocker.MagicMock())
 
     # THEN we should not add empty metadata
-    assert put_metadata_mock.call_count == 0
+    assert in_subsegment_mock.put_metadata.call_count == 0
 
 
-def test_tracer_method_does_not_add_empty_response_as_metadata(mocker, provider_stub):
+def test_tracer_method_does_not_add_empty_response_as_metadata(mocker, provider_stub, in_subsegment_mock):
     # GIVEN tracer is initialized
-    put_metadata_mock = mocker.MagicMock()
-    provider = provider_stub(put_metadata_mock=put_metadata_mock)
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
     tracer = Tracer(provider=provider)
 
     # WHEN capture_method decorator is used
@@ -212,7 +210,7 @@ def greeting(name, message):
     greeting(name="Foo", message="Bar")
 
     # THEN we should not add empty metadata
-    assert put_metadata_mock.call_count == 0
+    assert in_subsegment_mock.put_metadata.call_count == 0
 
 
 @mock.patch("aws_lambda_powertools.tracing.tracer.aws_xray_sdk.core.patch")
@@ -502,3 +500,37 @@ def generator_fn():
     assert put_metadata_mock_args["namespace"] == "booking"
     assert isinstance(put_metadata_mock_args["value"], ValueError)
     assert str(put_metadata_mock_args["value"]) == "test"
+
+
+def test_tracer_lambda_handler_does_not_add_response_as_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, auto_patch=False)
+
+    # WHEN capture_lambda_handler decorator is used
+    # and the handler response is empty
+    @tracer.capture_lambda_handler(capture_response=False)
+    def handler(event, context):
+        return "response"
+
+    handler({}, mocker.MagicMock())
+
+    # THEN we should not add any metadata
+    assert in_subsegment_mock.put_metadata.call_count == 0
+
+
+def test_tracer_method_does_not_add_response_as_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, auto_patch=False)
+
+    # WHEN capture_method decorator is used
+    # and the method response is empty
+    @tracer.capture_method(capture_response=False)
+    def greeting(name, message):
+        return "response"
+
+    greeting(name="Foo", message="Bar")
+
+    # THEN we should not add any metadata
+    assert in_subsegment_mock.put_metadata.call_count == 0