Merge pull request #128 from heitorlessa/feat/tracer-disallow-response-metadata

feat: option to disallow tracer to capture response as metadata

Author: heitorlessa

Diff for: 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

Diff for: 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"]

Diff for: 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