improv: override Tracer auto-capture response/exception via env vars (#259)

* improv: add capture_response/error env vars

* improv: add tests for capture_response/error env vars

* docs: update env vars

* improv: address Nicolas feedback; improve docs

* docs: fix wording on disable exception example

Author: heitorlessa

aws_lambda_powertools/shared/__init__.py

@@ -0,0 +1,4 @@
+import os
+
+TRACER_CAPTURE_RESPONSE_ENV: str = os.getenv("POWERTOOLS_TRACER_CAPTURE_RESPONSE", "true")
+TRACER_CAPTURE_ERROR_ENV: str = os.getenv("POWERTOOLS_TRACER_CAPTURE_ERROR", "true")

aws_lambda_powertools/shared/constants.py

@@ -0,0 +1,5 @@
+from distutils.util import strtobool
+
+
+def resolve_env_var_choice(env: str, choice: bool = None) -> bool:
+    return choice if choice is not None else strtobool(env)

aws_lambda_powertools/shared/functions.py

@@ -5,11 +5,14 @@
 import logging
 import os
 from distutils.util import strtobool
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import aws_xray_sdk
 import aws_xray_sdk.core
 
+from aws_lambda_powertools.shared.constants import TRACER_CAPTURE_ERROR_ENV, TRACER_CAPTURE_RESPONSE_ENV
+from aws_lambda_powertools.shared.functions import resolve_env_var_choice
+
 is_cold_start = True
 logger = logging.getLogger(__name__)
 
@@ -34,6 +37,10 @@ class Tracer:
         disable tracer (e.g. `"true", "True", "TRUE"`)
     POWERTOOLS_SERVICE_NAME : str
         service name
+    POWERTOOLS_TRACER_CAPTURE_RESPONSE : str
+        disable auto-capture response as metadata (e.g. `"true", "True", "TRUE"`)
+    POWERTOOLS_TRACER_CAPTURE_ERROR : str
+        disable auto-capture error as metadata (e.g. `"true", "True", "TRUE"`)
 
     Parameters
     ----------
@@ -226,7 +233,12 @@ 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, capture_response: bool = True):
+    def capture_lambda_handler(
+        self,
+        lambda_handler: Callable[[Dict, Any], Any] = None,
+        capture_response: Optional[bool] = None,
+        capture_error: Optional[bool] = None,
+    ):
         """Decorator to create subsegment for lambda handlers
 
         As Lambda follows (event, context) signature we can remove some of the boilerplate
@@ -237,7 +249,9 @@ def capture_lambda_handler(self, lambda_handler: Callable[[Dict, Any], Any] = No
         lambda_handler : Callable
             Method to annotate on
         capture_response : bool, optional
-            Instructs tracer to not include handler's response as metadata, by default True
+            Instructs tracer to not include handler's response as metadata
+        capture_error : bool, optional
+            Instructs tracer to not include handler's error as metadata, by default True
 
         Example
         -------
@@ -264,10 +278,15 @@ def handler(event, context):
         # 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)
+            return functools.partial(
+                self.capture_lambda_handler, capture_response=capture_response, capture_error=capture_error
+            )
 
         lambda_handler_name = lambda_handler.__name__
 
+        capture_response = resolve_env_var_choice(env=TRACER_CAPTURE_RESPONSE_ENV, choice=capture_response)
+        capture_error = resolve_env_var_choice(env=TRACER_CAPTURE_ERROR_ENV, choice=capture_error)
+
         @functools.wraps(lambda_handler)
         def decorate(event, context):
             with self.provider.in_subsegment(name=f"## {lambda_handler_name}") as subsegment:
@@ -290,15 +309,17 @@ def decorate(event, context):
                 except Exception as err:
                     logger.exception(f"Exception received from {lambda_handler_name}")
                     self._add_full_exception_as_metadata(
-                        method_name=lambda_handler_name, error=err, subsegment=subsegment
+                        method_name=lambda_handler_name, error=err, subsegment=subsegment, capture_error=capture_error
                     )
                     raise
 
                 return response
 
         return decorate
 
-    def capture_method(self, method: Callable = None, capture_response: bool = True):
+    def capture_method(
+        self, method: Callable = None, capture_response: Optional[bool] = None, capture_error: Optional[bool] = None
+    ):
         """Decorator to create subsegment for arbitrary functions
 
         It also captures both response and exceptions as metadata
@@ -317,7 +338,9 @@ def capture_method(self, method: Callable = None, capture_response: bool = True)
         method : Callable
             Method to annotate on
         capture_response : bool, optional
-            Instructs tracer to not include method's response as metadata, by default True
+            Instructs tracer to not include method's response as metadata
+        capture_error : bool, optional
+            Instructs tracer to not include handler's error as metadata, by default True
 
         Example
         -------
@@ -449,51 +472,65 @@ async def async_tasks():
         # 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)
+            return functools.partial(
+                self.capture_method, capture_response=capture_response, capture_error=capture_error
+            )
 
         method_name = f"{method.__name__}"
 
+        capture_response = resolve_env_var_choice(env=TRACER_CAPTURE_RESPONSE_ENV, choice=capture_response)
+        capture_error = resolve_env_var_choice(env=TRACER_CAPTURE_ERROR_ENV, choice=capture_error)
+
         if inspect.iscoroutinefunction(method):
             return self._decorate_async_function(
-                method=method, capture_response=capture_response, method_name=method_name
+                method=method, capture_response=capture_response, capture_error=capture_error, method_name=method_name
             )
         elif inspect.isgeneratorfunction(method):
             return self._decorate_generator_function(
-                method=method, capture_response=capture_response, method_name=method_name
+                method=method, capture_response=capture_response, capture_error=capture_error, method_name=method_name
             )
         elif hasattr(method, "__wrapped__") and inspect.isgeneratorfunction(method.__wrapped__):
             return self._decorate_generator_function_with_context_manager(
-                method=method, capture_response=capture_response, method_name=method_name
+                method=method, capture_response=capture_response, capture_error=capture_error, method_name=method_name
             )
         else:
             return self._decorate_sync_function(
-                method=method, capture_response=capture_response, method_name=method_name
+                method=method, capture_response=capture_response, capture_error=capture_error, method_name=method_name
             )
 
-    def _decorate_async_function(self, method: Callable = None, capture_response: bool = True, method_name: str = None):
+    def _decorate_async_function(
+        self,
+        method: Callable = None,
+        capture_response: Optional[bool] = None,
+        capture_error: Optional[bool] = None,
+        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(
-                        method_name=method_name,
-                        data=response,
-                        subsegment=subsegment,
-                        capture_response=capture_response,
+                        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(method_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(
+                        method_name=method_name, error=err, subsegment=subsegment, capture_error=capture_error
+                    )
                     raise
 
                 return response
 
         return decorate
 
     def _decorate_generator_function(
-        self, method: Callable = None, capture_response: bool = True, method_name: str = None
+        self,
+        method: Callable = None,
+        capture_response: Optional[bool] = None,
+        capture_error: Optional[bool] = None,
+        method_name: str = None,
     ):
         @functools.wraps(method)
         def decorate(*args, **kwargs):
@@ -506,15 +543,21 @@ def decorate(*args, **kwargs):
                     )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(
+                        method_name=method_name, error=err, subsegment=subsegment, capture_error=capture_error
+                    )
                     raise
 
                 return result
 
         return decorate
 
     def _decorate_generator_function_with_context_manager(
-        self, method: Callable = None, capture_response: bool = True, method_name: str = None
+        self,
+        method: Callable = None,
+        capture_response: Optional[bool] = None,
+        capture_error: Optional[bool] = None,
+        method_name: str = None,
     ):
         @functools.wraps(method)
         @contextlib.contextmanager
@@ -530,12 +573,20 @@ def decorate(*args, **kwargs):
                     )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(
+                        method_name=method_name, error=err, subsegment=subsegment, capture_error=capture_error
+                    )
                     raise
 
         return decorate
 
-    def _decorate_sync_function(self, method: Callable = None, capture_response: bool = True, method_name: str = None):
+    def _decorate_sync_function(
+        self,
+        method: Callable = None,
+        capture_response: Optional[bool] = None,
+        capture_error: Optional[bool] = None,
+        method_name: str = None,
+    ):
         @functools.wraps(method)
         def decorate(*args, **kwargs):
             with self.provider.in_subsegment(name=f"## {method_name}") as subsegment:
@@ -550,7 +601,9 @@ def decorate(*args, **kwargs):
                     )
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
-                    self._add_full_exception_as_metadata(method_name=method_name, error=err, subsegment=subsegment)
+                    self._add_full_exception_as_metadata(
+                        method_name=method_name, error=err, subsegment=subsegment, capture_error=capture_error
+                    )
                     raise
 
                 return response
@@ -562,7 +615,7 @@ def _add_response_as_metadata(
         method_name: str = None,
         data: Any = None,
         subsegment: aws_xray_sdk.core.models.subsegment = None,
-        capture_response: bool = True,
+        capture_response: Optional[bool] = None,
     ):
         """Add response as metadata for given subsegment
 
@@ -575,15 +628,19 @@ def _add_response_as_metadata(
         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
+            Do not include response as metadata
         """
         if data is None or not capture_response or subsegment is None:
             return
 
         subsegment.put_metadata(key=f"{method_name} response", value=data, namespace=self._config["service"])
 
     def _add_full_exception_as_metadata(
-        self, method_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,
+        capture_error: Optional[bool] = None,
     ):
         """Add full exception object as metadata for given subsegment
 
@@ -595,7 +652,12 @@ def _add_full_exception_as_metadata(
             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
+        capture_error : bool, optional
+            Do not include error as metadata, by default True
         """
+        if not capture_error:
+            return
+
         subsegment.put_metadata(key=f"{method_name} error", value=error, namespace=self._config["service"])
 
     @staticmethod

aws_lambda_powertools/tracing/tracer.py

@@ -65,12 +65,16 @@ def handler(event, context):
     ...
 ```
 
+### disabling response auto-capture
+
+> New in 1.9.0
+
 <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/>
 
+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.
 
 ```python:title=do_not_capture_response_as_metadata.py
 # Disables Tracer from capturing response and adding as metadata
@@ -80,6 +84,25 @@ def handler(event, context):
 	return "sensitive_information"
 ```
 
+### disabling exception auto-capture
+
+> New in 1.10.0
+
+<Note type="warning">
+  <strong>Can exceptions contain sensitive information from your Lambda handler or functions, where Tracer is used?</strong>
+  <br/><br/>
+</Note><br/>
+
+You can disable Tracer from capturing their exceptions as tracing metadata with <strong><code>capture_error=False</code></strong> parameter in both capture_lambda_handler and capture_method decorators.
+
+```python:title=do_not_capture_exception_as_metadata.py
+# Disables Tracer from capturing exception and adding as metadata
+# Useful when dealing with sensitive data
+@tracer.capture_lambda_handler(capture_error=False) # highlight-line
+def handler(event, context):
+	raise ValueError("some sensitive info in the stack trace...")
+```
+
 ### Annotations
 
 Annotations are key-values indexed by AWS X-Ray on a per trace basis. You can use them to filter traces as well as to create [Trace Groups](https://aws.amazon.com/about-aws/whats-new/2018/11/aws-xray-adds-the-ability-to-group-traces/).

docs/content/core/tracer.mdx

@@ -122,17 +122,23 @@ Utility | Description
 
 ## Environment variables
 
+<Note type="info">
+	<strong>Explicit parameters take precedence over environment variables.</strong><br/><br/>
+</Note>
+
 **Environment variables** used across suite of utilities.
 
-Environment variable | Description | Utility
-------------------------------------------------- | --------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- | -------------------------------------------------
-**POWERTOOLS_SERVICE_NAME** | Sets service name used for tracing namespace, metrics dimension and structured logging | All
-**POWERTOOLS_METRICS_NAMESPACE** | Sets namespace used for metrics | [Metrics](./core/metrics)
-**POWERTOOLS_TRACE_DISABLED** | Disables tracing | [Tracing](./core/tracer)
-**POWERTOOLS_TRACE_MIDDLEWARES** | Creates sub-segment for each custom middleware | [Middleware factory](./utilities/middleware_factory)
-**POWERTOOLS_LOGGER_LOG_EVENT** | Logs incoming event | [Logging](./core/logger)
-**POWERTOOLS_LOGGER_SAMPLE_RATE** | Debug log sampling | [Logging](./core/logger)
-**LOG_LEVEL** | Sets logging level | [Logging](./core/logger)
+Environment variable | Description | Utility | Default
+------------------------------------------------- | --------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- | ------------------------------------------------- | -------------------------------------------------
+**POWERTOOLS_SERVICE_NAME** | Sets service name used for tracing namespace, metrics dimension and structured logging | All | `"service_undefined"`
+**POWERTOOLS_METRICS_NAMESPACE** | Sets namespace used for metrics | [Metrics](./core/metrics) | `None`
+**POWERTOOLS_TRACE_DISABLED** | Disables tracing | [Tracing](./core/tracer) | `false`
+**POWERTOOLS_TRACER_CAPTURE_RESPONSE** | Captures Lambda or method return as metadata. | [Tracing](./core/tracer) | `true`
+**POWERTOOLS_TRACER_CAPTURE_ERROR** | Captures Lambda or method exception as metadata. | [Tracing](./core/tracer) | `true`
+**POWERTOOLS_TRACE_MIDDLEWARES** | Creates sub-segment for each custom middleware | [Middleware factory](./utilities/middleware_factory) | `false`
+**POWERTOOLS_LOGGER_LOG_EVENT** | Logs incoming event | [Logging](./core/logger) | `false`
+**POWERTOOLS_LOGGER_SAMPLE_RATE** | Debug log sampling | [Logging](./core/logger) | `0`
+**LOG_LEVEL** | Sets logging level | [Logging](./core/logger) | `INFO`
 
 ## Debug mode
 

docs/content/index.mdx

@@ -0,0 +1,9 @@
+from aws_lambda_powertools.shared.functions import resolve_env_var_choice
+
+
+def test_resolve_env_var_choice_explicit_wins_over_env_var():
+    assert resolve_env_var_choice(env="true", choice=False) is False
+
+
+def test_resolve_env_var_choice_env_wins_over_absent_explicit():
+    assert resolve_env_var_choice(env="true") == 1