Make enhanced metrics opt-in via an env var

Author: swfunc

datadog_lambda/metric.py

@@ -76,9 +76,18 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None):
         lambda_stats.distribution(metric_name, value, timestamp=timestamp, tags=tags)
 
 
+def are_enhanced_metrics_enabled():
+    """Check env var to find if enhanced metrics should be submitted
+    """
+    return os.environ.get("DD_ENHANCED_METRICS", "false").lower() == "true"
+
+
 def submit_invocations_metric(lambda_arn):
     """Increment aws.lambda.enhanced.invocations by 1
     """
+    if not are_enhanced_metrics_enabled():
+        return
+
     lambda_metric(
         "{}.invocations".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
         1,
@@ -89,6 +98,9 @@ def submit_invocations_metric(lambda_arn):
 def submit_errors_metric(lambda_arn):
     """Increment aws.lambda.enhanced.errors by 1
     """
+    if not are_enhanced_metrics_enabled():
+        return
+
     lambda_metric(
         "{}.errors".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
         1,

tests/test_wrapper.py

@@ -106,6 +106,8 @@ def lambda_handler(event, context):
         del os.environ["DD_LOGS_INJECTION"]
 
     def test_invocations_metric(self):
+        os.environ["DD_ENHANCED_METRICS"] = "True"
+
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
             lambda_metric("test.metric", 100)
@@ -129,7 +131,11 @@ def lambda_handler(event, context):
             ]
         )
 
+        del os.environ["DD_ENHANCED_METRICS"]
+
     def test_errors_metric(self):
+        os.environ["DD_ENHANCED_METRICS"] = "True"
+
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
             raise RuntimeError()
@@ -164,7 +170,11 @@ def lambda_handler(event, context):
             ]
         )
 
+        del os.environ["DD_ENHANCED_METRICS"]
+
     def test_enhanced_metrics_cold_start_tag(self):
+        os.environ["DD_ENHANCED_METRICS"] = "True"
+
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
             lambda_metric("test.metric", 100)
@@ -203,3 +213,18 @@ def lambda_handler(event, context):
                 ),
             ]
         )
+
+        del os.environ["DD_ENHANCED_METRICS"]
+
+    def test_no_enhanced_metrics_without_env_var(self):
+        @datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            raise RuntimeError()
+
+        lambda_event = {}
+
+        with self.assertRaises(RuntimeError):
+            lambda_handler(lambda_event, get_mock_context())
+
+        self.mock_wrapper_lambda_metric.assert_not_called()
+