fix(anomaly detection): get aggregation key from snuba data (#77498)

We can't assume that the aggregation key of a snuba timeseries data
point will be `"count"`, so we should set the aggregation key when we
see a data point with an aggregation value.

Author: mifu67

src/sentry/seer/anomaly_detection/utils.py

@@ -102,8 +102,21 @@ def format_historical_data(data: SnubaTSResult, dataset: Any) -> list[TimeSeries
             ts_point = TimeSeriesPoint(timestamp=date.timestamp(), value=count)
             formatted_data.append(ts_point)
     else:
+        # we don't know what the aggregation key of the query is
+        # so we should see it when we see a data point that has a value
+        agg_key = ""
         for datum in nested_data:
-            ts_point = TimeSeriesPoint(timestamp=datum.get("time"), value=datum.get("count", 0))
+            if len(datum) == 1:
+                # this data point has no value
+                ts_point = TimeSeriesPoint(timestamp=datum.get("time"), value=0)
+            else:
+                # if we don't know the aggregation key yet, we should set it
+                if not agg_key:
+                    for key in datum:  # only two keys in this dict
+                        if key != "time":
+                            agg_key = key
+                            break
+                ts_point = TimeSeriesPoint(timestamp=datum.get("time"), value=datum.get(agg_key, 0))
             formatted_data.append(ts_point)
     return formatted_data
 

tests/sentry/seer/anomaly_detection/test_store_data.py

@@ -58,6 +58,22 @@ def test_anomaly_detection_format_historical_data(self):
         result = format_historical_data(data, errors)
         assert result == expected_return_value
 
+    def test_anomaly_detection_format_historical_data_two(self):
+        """
+        Test a different aggregation key.
+        """
+        expected_return_value = [
+            {"timestamp": self.time_1_ts, "value": 0},
+            {"timestamp": self.time_2_ts, "value": 1},
+        ]
+        snuba_raw_data = [
+            {"time": self.time_1_ts},
+            {"count_unique_tags_sentry_user": 1, "time": self.time_2_ts},
+        ]
+        data = SnubaTSResult({"data": snuba_raw_data}, self.time_1_ts, self.time_2_ts, 3600)
+        result = format_historical_data(data, errors)
+        assert result == expected_return_value
+
     def test_anomaly_detection_fetch_historical_data(self):
         alert_rule = self.create_alert_rule(organization=self.organization, projects=[self.project])
         snuba_query = SnubaQuery.objects.get(id=alert_rule.snuba_query_id)