fix(metrics): support additional arguments in functions wrapped with log_metrics decorator (#3120)

Author: thegeorgeliu

aws_lambda_powertools/metrics/base.py

@@ -378,11 +378,11 @@ def handler(event, context):
             )
 
         @functools.wraps(lambda_handler)
-        def decorate(event, context):
+        def decorate(event, context, *args, **kwargs):
             try:
                 if default_dimensions:
                     self.set_default_dimensions(**default_dimensions)
-                response = lambda_handler(event, context)
+                response = lambda_handler(event, context, *args, **kwargs)
                 if capture_cold_start_metric:
                     self._add_cold_start_metric(context=context)
             finally:

aws_lambda_powertools/metrics/provider/base.py

@@ -199,9 +199,9 @@ def handler(event, context):
             )
 
         @functools.wraps(lambda_handler)
-        def decorate(event, context):
+        def decorate(event, context, *args, **kwargs):
             try:
-                response = lambda_handler(event, context)
+                response = lambda_handler(event, context, *args, **kwargs)
                 if capture_cold_start_metric:
                     self._add_cold_start_metric(context=context)
             finally:

tests/functional/metrics/test_metrics_cloudwatch_emf.py

@@ -355,6 +355,22 @@ def lambda_handler(evt, context):
     assert lambda_handler({}, {}) is True
 
 
+def test_log_metrics_decorator_with_additional_handler_args(namespace, service):
+    # GIVEN Metrics is initialized
+    my_metrics = Metrics(service=service, namespace=namespace)
+
+    # WHEN log_metrics is used to serialize metrics
+    # AND the wrapped function uses additional parameters
+    @my_metrics.log_metrics
+    def lambda_handler(evt, context, additional_arg, additional_kw_arg="default_value"):
+        return additional_arg, additional_kw_arg
+
+    # THEN the decorator should not raise any errors when
+    # the wrapped function is passed additional arguments
+    assert lambda_handler({}, {}, "arg_value", additional_kw_arg="kw_arg_value") == ("arg_value", "kw_arg_value")
+    assert lambda_handler({}, {}, "arg_value") == ("arg_value", "default_value")
+
+
 def test_schema_validation_incorrect_metric_resolution(metric, dimension):
     # GIVEN we pass a metric resolution that is not supported by CloudWatch
     metric["resolution"] = 10  # metric resolution must be 1 (High) or 60 (Standard)

tests/functional/metrics/test_metrics_datadog.py

@@ -136,6 +136,22 @@ def lambda_handler(evt, context):
         )
 
 
+def test_datadog_log_metrics_decorator_with_additional_handler_args():
+    # GIVEN DatadogMetrics is initialized
+    my_metrics = DatadogMetrics(flush_to_log=True)
+
+    # WHEN log_metrics is used to serialize metrics
+    # AND the wrapped function uses additional parameters
+    @my_metrics.log_metrics
+    def lambda_handler(evt, context, additional_arg, additional_kw_arg="default_value"):
+        return additional_arg, additional_kw_arg
+
+    # THEN the decorator should not raise any errors when
+    # the wrapped function is passed additional arguments
+    assert lambda_handler({}, {}, "arg_value", additional_kw_arg="kw_arg_value") == ("arg_value", "kw_arg_value")
+    assert lambda_handler({}, {}, "arg_value") == ("arg_value", "default_value")
+
+
 def test_metrics_with_default_namespace(capsys, namespace):
     # GIVEN DatadogMetrics is initialized with default namespace
     metrics = DatadogMetrics(flush_to_log=True)

tests/functional/metrics/test_metrics_provider.py

@@ -60,3 +60,19 @@ def lambda_handler(evt, context):
     # THEN log_metrics should invoke the function it decorates
     # and return no error if we have a namespace and dimension
     assert lambda_handler({}, {}) is True
+
+
+def test_metrics_provider_class_decorator_with_additional_handler_args():
+    # GIVEN Metrics is initialized
+    my_metrics = Metrics()
+
+    # WHEN log_metrics is used to serialize metrics
+    # AND the wrapped function uses additional parameters
+    @my_metrics.log_metrics
+    def lambda_handler(evt, context, additional_arg, additional_kw_arg="default_value"):
+        return additional_arg, additional_kw_arg
+
+    # THEN the decorator should not raise any errors when
+    # the wrapped function is passed additional arguments
+    assert lambda_handler({}, {}, "arg_value", additional_kw_arg="kw_arg_value") == ("arg_value", "kw_arg_value")
+    assert lambda_handler({}, {}, "arg_value") == ("arg_value", "default_value")