feat: send metrics with timestamps to the extension

This is now supported by both the go extension and bottlecap, so we can
safely send metrics to the extension and not have to ship the ones with
timestamps directly to the Datadog API.

Author: apiarian-datadog

datadog_lambda/metric.py

@@ -76,39 +76,25 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
     tags = [] if tags is None else list(tags)
     tags.append(dd_lambda_layer_tag)
 
-    if should_use_extension and timestamp is not None:
-        # The extension does not support timestamps for distributions so we create a
-        # a thread stats writer to submit metrics with timestamps to the API
-        timestamp_ceiling = int(
-            (datetime.now() - timedelta(hours=4)).timestamp()
-        )  # 4 hours ago
-        if isinstance(timestamp, datetime):
-            timestamp = int(timestamp.timestamp())
-        if timestamp_ceiling > timestamp:
-            logger.warning(
-                "Timestamp %s is older than 4 hours, not submitting metric %s",
-                timestamp,
-                metric_name,
-            )
-            return
-        global extension_thread_stats
-        if extension_thread_stats is None:
-            from datadog_lambda.api import init_api
-            from datadog_lambda.thread_stats_writer import ThreadStatsWriter
-
-            init_api()
-            extension_thread_stats = ThreadStatsWriter(flush_in_thread)
-
-        extension_thread_stats.distribution(
-            metric_name, value, tags=tags, timestamp=timestamp
-        )
-        return
-
     if should_use_extension:
+        if timestamp is not None:
+            if isinstance(timestamp, datetime):
+                timestamp = int(timestamp.timestamp())
+
+            timestamp_floor = int((datetime.now() - timedelta(hours=4)).timestamp())
+            if timestamp < timestamp_floor:
+                logger.warning(
+                    "Timestamp %s is older than 4 hours, not submitting metric %s",
+                    timestamp,
+                    metric_name,
+                )
+                return
+
         logger.debug(
             "Sending metric %s value %s to Datadog via extension", metric_name, value
         )
         lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
+
     else:
         if flush_to_logs or force_async:
             write_metric_point_to_stdout(

tests/test_metric.py

@@ -1,16 +1,15 @@
 import os
 import unittest
-
-from unittest.mock import patch, call
+from datetime import datetime, timedelta
+from unittest.mock import call, patch
 
 from botocore.exceptions import ClientError as BotocoreClientError
 from datadog.api.exceptions import ClientError
-from datetime import datetime, timedelta
 
-from datadog_lambda.metric import lambda_metric, flush_stats
-from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY
-from datadog_lambda.thread_stats_writer import ThreadStatsWriter
+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.tags import dd_lambda_layer_tag
+from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 
 
 class TestLambdaMetric(unittest.TestCase):
@@ -53,22 +52,31 @@ def test_lambda_metric_timestamp_with_extension(self):
         timestamp = int((datetime.now() - delta).timestamp())
 
         lambda_metric("test_timestamp", 1, timestamp)
-        self.mock_metric_lambda_stats.distribution.assert_not_called()
-        self.mock_metric_extension_thread_stats.distribution.assert_called_with(
-            "test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag]
+        self.mock_metric_lambda_stats.distribution.assert_has_calls(
+            [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])]
         )
+        self.mock_metric_extension_thread_stats.distribution.assert_not_called()
 
     @patch("datadog_lambda.metric.should_use_extension", True)
     def test_lambda_metric_datetime_with_extension(self):
         patcher = patch("datadog_lambda.metric.extension_thread_stats")
         self.mock_metric_extension_thread_stats = patcher.start()
         self.addCleanup(patcher.stop)
 
-        delta = timedelta(hours=5)
+        delta = timedelta(minutes=1)
         timestamp = datetime.now() - delta
 
-        lambda_metric("test_timestamp", 1, timestamp)
-        self.mock_metric_lambda_stats.distribution.assert_not_called()
+        lambda_metric("test_datetime_timestamp", 0, timestamp)
+        self.mock_metric_lambda_stats.distribution.assert_has_calls(
+            [
+                call(
+                    "test_datetime_timestamp",
+                    0,
+                    timestamp=int(timestamp.timestamp()),
+                    tags=[dd_lambda_layer_tag],
+                )
+            ]
+        )
         self.mock_metric_extension_thread_stats.distribution.assert_not_called()
 
     @patch("datadog_lambda.metric.should_use_extension", True)