Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file.
4 changes: 4 additions & 0 deletions aws_lambda_powertools/shared/constants.py
Original file line number Diff line number Diff line change
@@ -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")
5 changes: 5 additions & 0 deletions aws_lambda_powertools/shared/functions.py
Original file line number Diff line number Diff line change
@@ -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)
116 changes: 89 additions & 27 deletions aws_lambda_powertools/tracing/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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
----------
Expand Down Expand Up @@ -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
Expand All @@ -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
-------
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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
-------
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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
Expand Down
25 changes: 24 additions & 1 deletion docs/content/core/tracer.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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/).
Expand Down
24 changes: 15 additions & 9 deletions docs/content/index.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -116,17 +116,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

Expand Down
9 changes: 9 additions & 0 deletions tests/functional/test_shared_functions.py
Original file line number Diff line number Diff line change
@@ -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
Loading