refs(alert_rules): Switch alert rule create/update code to use aggregate. (#18847)

This moves all code to use the new `SnubaQuery.aggregate` field, and pass `aggregate` around
(almost) everywhere. We still need to update this in the api, and there are still some places that
need to change to be able to support aggregates that don't map to our previous default aggregates/

bin/load-mocks

@@ -663,7 +663,7 @@ def main(num_events=1, extra_events=False):
             try:
                 # Metric alerts
                 alert_rule = create_alert_rule(
-                    org, [project], "My Alert Rule", "level:error", QueryAggregations.TOTAL, 10, 1
+                    org, [project], "My Alert Rule", "level:error", "count()", 10, 1
                 )
                 create_alert_rule_trigger(alert_rule, "critical", AlertRuleThresholdType.ABOVE, 10)
                 create_incident(

src/sentry/incidents/endpoints/serializers.py

@@ -34,6 +34,7 @@ from sentry.models.organizationmember import OrganizationMember
 from sentry.models.team import Team
 from sentry.models.user import User
 from sentry.snuba.models import QueryAggregations
+from sentry.snuba.subscriptions import aggregation_function_translations
 from sentry.utils.compat import zip
 
 
@@ -279,7 +280,7 @@ class AlertRuleSerializer(CamelSnakeModelSerializer):
         required=True, min_value=1, max_value=int(timedelta(days=1).total_seconds() / 60)
     )
     threshold_period = serializers.IntegerField(default=1, min_value=1, max_value=20)
-    aggregation = serializers.IntegerField(required=False)
+    aggregation = serializers.IntegerField(required=True)
 
     class Meta:
         model = AlertRule
@@ -314,6 +315,8 @@ class AlertRuleSerializer(CamelSnakeModelSerializer):
         This includes ensuring there is either 1 or 2 triggers, which each have actions, and have proper thresholds set.
         The critical trigger should both alert and resolve 'after' the warning trigger (whether that means > or < the value depends on threshold type).
         """
+        data["aggregate"] = aggregation_function_translations[data.pop("aggregation")]
+
         triggers = data.get("triggers", [])
         if not triggers:
             raise serializers.ValidationError("Must include at least one trigger")

src/sentry/incidents/logic.py

@@ -34,7 +34,6 @@ from sentry.incidents.models import (
 from sentry.models import Integration, Project
 from sentry.snuba.models import QueryDatasets
 from sentry.snuba.subscriptions import (
-    aggregate_to_query_aggregation,
     bulk_create_snuba_subscriptions,
     bulk_delete_snuba_subscriptions,
     create_snuba_query,
@@ -493,7 +492,7 @@ def create_alert_rule(
     projects,
     name,
     query,
-    aggregation,
+    aggregate,
     time_window,
     threshold_period,
     environment=None,
@@ -509,7 +508,7 @@ def create_alert_rule(
     :param name: Name for the alert rule. This will be used as part of the
     incident name, and must be unique per project
     :param query: An event search query to subscribe to and monitor for alerts
-    :param aggregation: A QueryAggregation to fetch for this alert rule
+    :param aggregate: A string representing the aggregate used in this alert rule
     :param time_window: Time period to aggregate over, in minutes
     :param environment: An optional environment that this rule applies to
     :param threshold_period: How many update periods the value of the
@@ -530,7 +529,7 @@ def create_alert_rule(
         snuba_query = create_snuba_query(
             dataset,
             query,
-            aggregation,
+            aggregate,
             timedelta(minutes=time_window),
             timedelta(minutes=resolution),
             environment,
@@ -597,7 +596,7 @@ def update_alert_rule(
     projects=None,
     name=None,
     query=None,
-    aggregation=None,
+    aggregate=None,
     time_window=None,
     environment=None,
     threshold_period=None,
@@ -613,7 +612,7 @@ def update_alert_rule(
     :param name: Name for the alert rule. This will be used as part of the
     incident name, and must be unique per project.
     :param query: An event search query to subscribe to and monitor for alerts
-    :param aggregation: An AlertRuleAggregation that we want to fetch for this alert rule
+    :param aggregate: A string representing the aggregate used in this alert rule
     :param time_window: Time period to aggregate over, in minutes.
     :param environment: An optional environment that this rule applies to
     :param threshold_period: How many update periods the value of the
@@ -638,8 +637,8 @@ def update_alert_rule(
     if query is not None:
         validate_alert_rule_query(query)
         updated_query_fields["query"] = query
-    if aggregation is not None:
-        updated_query_fields["aggregation"] = aggregation
+    if aggregate is not None:
+        updated_query_fields["aggregate"] = aggregate
     if time_window:
         updated_query_fields["time_window"] = timedelta(minutes=time_window)
     if threshold_period:
@@ -658,9 +657,7 @@ def update_alert_rule(
             updated_query_fields.setdefault("query", snuba_query.query)
             # XXX: We use the alert rule aggregation here since currently we're
             # expecting the enum value to be passed.
-            updated_query_fields.setdefault(
-                "aggregation", aggregate_to_query_aggregation[snuba_query.aggregate]
-            )
+            updated_query_fields.setdefault("aggregate", snuba_query.aggregate)
             updated_query_fields.setdefault(
                 "time_window", timedelta(seconds=snuba_query.time_window)
             )
@@ -674,7 +671,7 @@ def update_alert_rule(
         existing_subs = []
         if (
             query is not None
-            or aggregation is not None
+            or aggregate is not None
             or time_window is not None
             or projects is not None
             or include_all_projects is not None

src/sentry/snuba/subscriptions.py

@@ -22,17 +22,7 @@ aggregate_to_query_aggregation = {
 }
 
 
-def translate_aggregation(aggregation):
-    """
-    Temporary function to translate `QueryAggregations` into the discover aggregation
-    function format
-    :param aggregation:
-    :return: A string representing the aggregate function
-    """
-    return aggregation_function_translations[aggregation]
-
-
-def create_snuba_query(dataset, query, aggregation, time_window, resolution, environment):
+def create_snuba_query(dataset, query, aggregate, time_window, resolution, environment):
     """
     Creates a SnubaQuery.
 
@@ -48,14 +38,14 @@ def create_snuba_query(dataset, query, aggregation, time_window, resolution, env
     return SnubaQuery.objects.create(
         dataset=dataset.value,
         query=query,
-        aggregate=translate_aggregation(aggregation),
+        aggregate=aggregate,
         time_window=int(time_window.total_seconds()),
         resolution=int(resolution.total_seconds()),
         environment=environment,
     )
 
 
-def update_snuba_query(snuba_query, query, aggregation, time_window, resolution, environment):
+def update_snuba_query(snuba_query, query, aggregate, time_window, resolution, environment):
     """
     Updates a SnubaQuery. Triggers updates to any related QuerySubscriptions.
 
@@ -63,7 +53,7 @@ def update_snuba_query(snuba_query, query, aggregation, time_window, resolution,
     :param dataset: The snuba dataset to query and aggregate over
     :param query: An event search query that we can parse and convert into a
     set of Snuba conditions
-    :param aggregation: An aggregation to calculate over the time window
+    :param aggregate: An aggregate to calculate over the time window
     :param time_window: The time window to aggregate over
     :param resolution: How often to receive updates/bucket size
     :param environment: An optional environment to filter by
@@ -73,7 +63,7 @@ def update_snuba_query(snuba_query, query, aggregation, time_window, resolution,
         query_subscriptions = list(snuba_query.subscriptions.all())
         snuba_query.update(
             query=query,
-            aggregate=translate_aggregation(aggregation),
+            aggregate=aggregate,
             time_window=int(time_window.total_seconds()),
             resolution=int(resolution.total_seconds()),
             environment=environment,

src/sentry/testutils/factories.py

@@ -73,7 +73,6 @@ from sentry.models import (
 )
 from sentry.models.integrationfeature import Feature, IntegrationFeature
 from sentry.signals import project_created
-from sentry.snuba.models import QueryAggregations
 from sentry.utils import loremipsum, json
 
 
@@ -336,7 +335,7 @@ class Factories(object):
                 "workspace": integration_id,
                 "channel_id": channel_id or "123453",
                 "channel": channel_name or "#general",
-            },
+            }
         ]
         return Factories.create_project_rule(project, action_data)
 
@@ -858,7 +857,7 @@ class Factories(object):
         projects,
         name=None,
         query="level:error",
-        aggregation=QueryAggregations.TOTAL,
+        aggregate="count()",
         time_window=10,
         threshold_period=1,
         include_all_projects=False,
@@ -874,7 +873,7 @@ class Factories(object):
             projects,
             name,
             query,
-            aggregation,
+            aggregate,
             time_window,
             threshold_period,
             environment=environment,

tests/acceptance/test_incidents.py

@@ -28,13 +28,7 @@ class OrganizationIncidentsListTest(AcceptanceTestCase, SnubaTestCase):
 
     def test_incidents_list(self):
         alert_rule = create_alert_rule(
-            self.organization,
-            [self.project],
-            "hello",
-            "level:error",
-            QueryAggregations.TOTAL,
-            10,
-            1,
+            self.organization, [self.project], "hello", "level:error", "count()", 10, 1
         )
 
         incident = create_incident(

tests/sentry/api/serializers/test_alert_rule.py

@@ -12,7 +12,6 @@ from sentry.api.serializers.models.alert_rule import (
 from sentry.models import Rule
 from sentry.incidents.logic import create_alert_rule, create_alert_rule_trigger
 from sentry.incidents.models import AlertRuleThresholdType
-from sentry.snuba.models import QueryAggregations
 from sentry.snuba.subscriptions import aggregate_to_query_aggregation
 from sentry.testutils import TestCase, APITestCase
 
@@ -77,13 +76,7 @@ class BaseAlertRuleSerializerTest(object):
 class AlertRuleSerializerTest(BaseAlertRuleSerializerTest, TestCase):
     def test_simple(self):
         alert_rule = create_alert_rule(
-            self.organization,
-            [self.project],
-            "hello",
-            "level:error",
-            QueryAggregations.TOTAL,
-            10,
-            1,
+            self.organization, [self.project], "hello", "level:error", "count()", 10, 1
         )
         result = serialize(alert_rule)
         self.assert_alert_rule_serialized(alert_rule, result)

tests/sentry/incidents/endpoints/test_organization_alert_rule_index.py

@@ -6,11 +6,8 @@ from freezegun import freeze_time
 from sentry.api.serializers import serialize
 from sentry.incidents.logic import create_alert_rule
 from sentry.incidents.models import AlertRule
-from sentry.snuba.models import QueryAggregations
 from sentry.testutils import APITestCase
 
-# from sentry.incidents.endpoints.serializers import CRITICAL_TRIGGER_LABEL, WARNING_TRIGGER_LABEL
-
 
 class AlertRuleListEndpointTest(APITestCase):
     endpoint = "sentry-api-0-organization-alert-rules"
@@ -30,13 +27,7 @@ class AlertRuleListEndpointTest(APITestCase):
     def test_simple(self):
         self.create_team(organization=self.organization, members=[self.user])
         alert_rule = create_alert_rule(
-            self.organization,
-            [self.project],
-            "hello",
-            "level:error",
-            QueryAggregations.TOTAL,
-            10,
-            1,
+            self.organization, [self.project], "hello", "level:error", "count()", 10, 1
         )
 
         self.login_as(self.user)

tests/sentry/incidents/endpoints/test_project_alert_rule_details.py

@@ -5,8 +5,6 @@ from exam import fixture
 from sentry.api.serializers import serialize
 from sentry.incidents.logic import create_alert_rule
 from sentry.incidents.models import AlertRule, AlertRuleStatus, Incident, IncidentStatus
-from sentry.snuba.models import QueryAggregations
-from sentry.snuba.subscriptions import aggregate_to_query_aggregation
 from sentry.testutils import APITestCase
 
 
@@ -77,13 +75,7 @@ class AlertRuleDetailsBase(object):
     @fixture
     def alert_rule(self):
         return create_alert_rule(
-            self.organization,
-            [self.project],
-            "hello",
-            "level:error",
-            QueryAggregations.TOTAL,
-            10,
-            1,
+            self.organization, [self.project], "hello", "level:error", "count()", 10, 1
         )
 
     def test_invalid_rule_id(self):
@@ -148,13 +140,7 @@ class AlertRuleDetailsPutEndpointTest(AlertRuleDetailsBase, APITestCase):
 
     def test_not_updated_fields(self):
         test_params = self.valid_params.copy()
-        test_params.update(
-            {
-                "aggregation": aggregate_to_query_aggregation[
-                    self.alert_rule.snuba_query.aggregate
-                ].value
-            }
-        )
+        test_params["aggregate"] = self.alert_rule.snuba_query.aggregate
 
         self.create_member(
             user=self.user, organization=self.organization, role="owner", teams=[self.team]

tests/sentry/incidents/endpoints/test_project_alert_rule_index.py

@@ -11,7 +11,6 @@ from freezegun import freeze_time
 from sentry.api.serializers import serialize
 from sentry.incidents.logic import create_alert_rule
 from sentry.incidents.models import AlertRule
-from sentry.snuba.models import QueryAggregations
 from sentry.testutils.helpers.datetime import before_now
 from sentry.testutils import TestCase, APITestCase
 from tests.sentry.api.serializers.test_alert_rule import BaseAlertRuleSerializerTest
@@ -38,13 +37,7 @@ class AlertRuleListEndpointTest(APITestCase):
     def test_simple(self):
         self.create_team(organization=self.organization, members=[self.user])
         alert_rule = create_alert_rule(
-            self.organization,
-            [self.project],
-            "hello",
-            "level:error",
-            QueryAggregations.TOTAL,
-            10,
-            1,
+            self.organization, [self.project], "hello", "level:error", "count()", 10, 1
         )
 
         self.login_as(self.user)