improv: override Tracer auto-capture response/exception via env vars

aws_lambda_powertools/shared/constants.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/functions.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/tracing/tracer.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

docs/content/index.mdx

@@ -118,15 +118,16 @@ Utility | Description
 
 **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** | Overrides tracer from auto-capturing response 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
 

tests/functional/test_shared_functions.py

@@ -0,0 +1,11 @@
+import os
+
+from aws_lambda_powertools.shared.functions import resolve_env_var_choice
+
+
+def test_explicit_wins_over_env_var():
+    choice_env = os.getenv("CHOICE", True)
+
+    choice = resolve_env_var_choice(env=choice_env, choice=False)
+
+    assert choice is False

tests/unit/test_tracing.py

@@ -502,13 +502,15 @@ def generator_fn():
     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):
+def test_tracer_lambda_handler_override_response_as_metadata(mocker, provider_stub, in_subsegment_mock):
     # GIVEN tracer is initialized
     provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+
+    mocker.patch("aws_lambda_powertools.tracing.tracer.TRACER_CAPTURE_RESPONSE_ENV", return_value=True)
     tracer = Tracer(provider=provider, auto_patch=False)
 
     # WHEN capture_lambda_handler decorator is used
-    # and the handler response is empty
+    # with capture_response set to False
     @tracer.capture_lambda_handler(capture_response=False)
     def handler(event, context):
         return "response"
@@ -519,7 +521,7 @@ def handler(event, context):
     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):
+def test_tracer_method_override_response_as_metadata(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)