From 5f4c650812170daff7a5c827e8f22b763cbc2a75 Mon Sep 17 00:00:00 2001
From: "sentry-junior[bot]"
 <264270552+sentry-junior[bot]@users.noreply.github.com>
Date: Fri, 23 Oct 2026 14:17:06 +0000
Subject: [PATCH] fix(fetcher): Tolerate missing unit key in
 warn_multiple_units

Datadog can omit the unit key from a series entirely instead of returning null. Treat a missing key the same as null so the cronjob no longer crashes with KeyError.

Refs PRODENG-2990

Co-Authored-By: Trent Schmidt <trent.schmidt@sentry.io>
---
 py/tests/test_fetcher.py              | 27 +++++++++++++++++++++++++++
 py/usageaccountant/datadog_fetcher.py |  2 +-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/py/tests/test_fetcher.py b/py/tests/test_fetcher.py
index 73564c3..970a4fd 100644
--- a/py/tests/test_fetcher.py
+++ b/py/tests/test_fetcher.py
@@ -281,6 +281,33 @@ def test_main(self, mock_time: Mock, mock_query_dd: Mock) -> None:
         msg2 = self.broker.consume(Partition(self.topic, 0), 1)
         assert msg2 is None
 
+    @patch("usageaccountant.datadog_fetcher.query_datadog")
+    @patch("time.time")
+    def test_main_series_missing_unit_key(
+        self, mock_time: Mock, mock_query_dd: Mock
+    ) -> None:
+        """
+        Datadog can omit the "unit" key from a series entirely. With a
+        configured unit, main() should still record the data.
+        """
+        self.good_response["series"][0].pop("unit")  # type: ignore
+        mock_query_dd.return_value = self.good_response
+        mock_time.return_value = 1
+
+        ddf.main(
+            query_file=StringIO(self.query_str),
+            start_time=1,
+            period_seconds=2,
+            usage_accumulator=self.usage_accumulator,
+            dry_run=False,
+        )
+
+        msg1 = self.broker.consume(Partition(self.topic, 0), 0)
+        assert msg1 is not None
+
+        msg2 = self.broker.consume(Partition(self.topic, 0), 1)
+        assert msg2 is None
+
     @patch("usageaccountant.datadog_fetcher.query_datadog")
     @patch("time.time")
     def test_main_no_unit(self, mock_time: Mock, mock_query_dd: Mock) -> None:
diff --git a/py/usageaccountant/datadog_fetcher.py b/py/usageaccountant/datadog_fetcher.py
index 6400f64..af509a8 100644
--- a/py/usageaccountant/datadog_fetcher.py
+++ b/py/usageaccountant/datadog_fetcher.py
@@ -174,7 +174,7 @@ def warn_multiple_units(series_list: Sequence[DatadogResponseSeries]) -> None:
     if not series_list:
         return
 
-    unit_list = series_list[0]["unit"]
+    unit_list = series_list[0].get("unit")
     if unit_list and unit_list[1] is not None:
         logger.warning(f"Received multiple units from Datadog: {unit_list}.")