feat: add a more explicit MetricsHandler process for routing metrics

Author: apiarian-datadog

datadog_lambda/metric.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import enum
 import logging
 import os
 import time
@@ -15,25 +16,44 @@
 
 logger = logging.getLogger(__name__)
 
-lambda_stats = None
 
-flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
+class MetricsHandler(enum.Enum):
+    EXTENSION = "extension"
+    FORWARDER = "forwarder"
+    DATADOG_API = "datadog_api"
+
+
+def _select_metrics_handler():
+    if should_use_extension:
+        return MetricsHandler.EXTENSION
+    if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+        return MetricsHandler.FORWARDER
+    return MetricsHandler.DATADOG_API
+
 
-if should_use_extension:
+metrics_handler = _select_metrics_handler()
+logger.debug("identified primary metrics handler as %s", metrics_handler)
+
+
+lambda_stats = None
+if metrics_handler == MetricsHandler.EXTENSION:
     from datadog_lambda.statsd_writer import StatsDWriter
 
     lambda_stats = StatsDWriter()
-else:
+
+elif metrics_handler == MetricsHandler.DATADOG_API:
     # Periodical flushing in a background thread is NOT guaranteed to succeed
     # and leads to data loss. When disabled, metrics are only flushed at the
     # end of invocation. To make metrics submitted from a long-running Lambda
     # function available sooner, consider using the Datadog Lambda extension.
     from datadog_lambda.api import init_api
     from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 
+    flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
     init_api()
     lambda_stats = ThreadStatsWriter(flush_in_thread)
 
+
 enhanced_metrics_enabled = (
     os.environ.get("DD_ENHANCED_METRICS", "true").lower() == "true"
 )
@@ -44,16 +64,19 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
     Submit a data point to Datadog distribution metrics.
     https://docs.datadoghq.com/graphing/metrics/distributions/
 
-    When DD_FLUSH_TO_LOG is True, write metric to log, and
-    wait for the Datadog Log Forwarder Lambda function to submit
-    the metrics asynchronously.
+    If the Datadog Lambda Extension is present, metrics are submitted to its
+    dogstatsd endpoint.
+
+    When DD_FLUSH_TO_LOG is True or force_async is True, write metric to log,
+    and wait for the Datadog Log Forwarder Lambda function to submit the
+    metrics asynchronously.
 
     Otherwise, the metrics will be submitted to the Datadog API
     periodically and at the end of the function execution in a
     background thread.
 
-    Note that if the extension is present, it will override the DD_FLUSH_TO_LOG value
-    and always use the layer to send metrics to the extension
+    Note that if the extension is present, it will override the DD_FLUSH_TO_LOG
+    value and always use the layer to send metrics to the extension
     """
     if not metric_name or not isinstance(metric_name, str):
         logger.warning(
@@ -71,11 +94,10 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
         )
         return
 
-    flush_to_logs = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
     tags = [] if tags is None else list(tags)
     tags.append(dd_lambda_layer_tag)
 
-    if should_use_extension:
+    if metrics_handler == MetricsHandler.EXTENSION:
         if timestamp is not None:
             if isinstance(timestamp, datetime):
                 timestamp = int(timestamp.timestamp())
@@ -94,15 +116,20 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
         )
         lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
 
+    elif force_async or (metrics_handler == MetricsHandler.FORWARDER):
+        write_metric_point_to_stdout(metric_name, value, timestamp=timestamp, tags=tags)
+
+    elif metrics_handler == MetricsHandler.DATADOG_API:
+        lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
+
     else:
-        if flush_to_logs or force_async:
-            write_metric_point_to_stdout(
-                metric_name, value, timestamp=timestamp, tags=tags
-            )
-        else:
-            lambda_stats.distribution(
-                metric_name, value, tags=tags, timestamp=timestamp
-            )
+        # This should be qutie impossible, but let's at least log a message if
+        # it somehow happens.
+        logger.debug(
+            "Metric %s cannot be submitted because the metrics handler is not configured: %s",
+            metric_name,
+            metrics_handler,
+        )
 
 
 def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=None):

tests/test_metric.py

@@ -7,7 +7,12 @@
 from datadog.api.exceptions import ClientError
 
 from datadog_lambda.api import KMS_ENCRYPTION_CONTEXT_KEY, decrypt_kms_api_key
-from datadog_lambda.metric import flush_stats, lambda_metric
+from datadog_lambda.metric import (
+    MetricsHandler,
+    _select_metrics_handler,
+    flush_stats,
+    lambda_metric,
+)
 from datadog_lambda.tags import dd_lambda_layer_tag
 from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 
@@ -34,15 +39,31 @@ def test_lambda_metric_tagged_with_dd_lambda_layer(self):
 
     # let's fake that the extension is present, this should override DD_FLUSH_TO_LOG
     @patch("datadog_lambda.metric.should_use_extension", True)
-    def test_lambda_metric_flush_to_log_with_extension(self):
+    def test_select_metrics_handler_extension_despite_flush_to_logs(self):
         os.environ["DD_FLUSH_TO_LOG"] = "True"
+        self.assertEqual(MetricsHandler.EXTENSION, _select_metrics_handler())
+        del os.environ["DD_FLUSH_TO_LOG"]
+
+    @patch("datadog_lambda.metric.should_use_extension", False)
+    def test_select_metrics_handler_forwarder_when_flush_to_logs(self):
+        os.environ["DD_FLUSH_TO_LOG"] = "True"
+        self.assertEqual(MetricsHandler.FORWARDER, _select_metrics_handler())
+        del os.environ["DD_FLUSH_TO_LOG"]
+
+    @patch("datadog_lambda.metric.should_use_extension", False)
+    def test_select_metrics_handler_dd_api_fallback(self):
+        os.environ["DD_FLUSH_TO_LOG"] = "False"
+        self.assertEqual(MetricsHandler.DATADOG_API, _select_metrics_handler())
+        del os.environ["DD_FLUSH_TO_LOG"]
+
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
+    def test_lambda_metric_flush_to_log_with_extension(self):
         lambda_metric("test", 1)
         self.mock_metric_lambda_stats.distribution.assert_has_calls(
             [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])]
         )
-        del os.environ["DD_FLUSH_TO_LOG"]
 
-    @patch("datadog_lambda.metric.should_use_extension", True)
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_timestamp_with_extension(self):
         delta = timedelta(minutes=1)
         timestamp = int((datetime.now() - delta).timestamp())
@@ -52,7 +73,7 @@ def test_lambda_metric_timestamp_with_extension(self):
             [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])]
         )
 
-    @patch("datadog_lambda.metric.should_use_extension", True)
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_datetime_with_extension(self):
         delta = timedelta(minutes=1)
         timestamp = datetime.now() - delta
@@ -69,22 +90,19 @@ def test_lambda_metric_datetime_with_extension(self):
             ]
         )
 
-    @patch("datadog_lambda.metric.should_use_extension", True)
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_invalid_timestamp_with_extension(self):
         delta = timedelta(hours=5)
         timestamp = int((datetime.now() - delta).timestamp())
 
         lambda_metric("test_timestamp", 1, timestamp)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
 
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.FORWARDER)
     def test_lambda_metric_flush_to_log(self):
-        os.environ["DD_FLUSH_TO_LOG"] = "True"
-
         lambda_metric("test", 1)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
 
-        del os.environ["DD_FLUSH_TO_LOG"]
-
     @patch("datadog_lambda.metric.logger.warning")
     def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning):
         lambda_metric(None, 1)