Skip to content

Commit

Permalink
improv: add support for custom lambda handlers with kwargs #242 (#269)
Browse files Browse the repository at this point in the history
* improv: support kwargs

* chore: add a comment to not forget about async tracer
  • Loading branch information
heitorlessa authored Jan 17, 2021
1 parent 454d669 commit 1d0ee0e
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
5 changes: 3 additions & 2 deletions aws_lambda_powertools/tracing/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,7 @@ def handler(event, context):
)

@functools.wraps(lambda_handler)
def decorate(event, context):
def decorate(event, context, **kwargs):
with self.provider.in_subsegment(name=f"## {lambda_handler_name}") as subsegment:
global is_cold_start
if is_cold_start:
Expand All @@ -300,7 +300,7 @@ def decorate(event, context):

try:
logger.debug("Calling lambda handler")
response = lambda_handler(event, context)
response = lambda_handler(event, context, **kwargs)
logger.debug("Received lambda handler response successfully")
self._add_response_as_metadata(
method_name=lambda_handler_name,
Expand Down Expand Up @@ -487,6 +487,7 @@ async def async_tasks():
env=os.getenv(constants.TRACER_CAPTURE_ERROR_ENV, "true"), choice=capture_error
)

# Maintenance: Need a factory/builder here to simplify this now
if inspect.iscoroutinefunction(method):
return self._decorate_async_function(
method=method, capture_response=capture_response, capture_error=capture_error, method_name=method_name
Expand Down
13 changes: 13 additions & 0 deletions tests/functional/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,19 @@ def handler(event, context):
handler({}, {})


def test_capture_lambda_handler_with_additional_kwargs(dummy_response):
# GIVEN tracer lambda handler decorator is used
tracer = Tracer(disabled=True)

# WHEN a lambda handler signature has additional keyword arguments
@tracer.capture_lambda_handler
def handler(event, context, my_extra_option=None, **kwargs):
return dummy_response

# THEN tracer should not raise an Exception
handler({}, {}, blah="blah")


def test_capture_method(dummy_response):
# GIVEN tracer method decorator is used
tracer = Tracer(disabled=True)
Expand Down

0 comments on commit 1d0ee0e

Please sign in to comment.