feat(upsampling) - Organization Events API error upsampling support (#95473)

Support projects with error upsampling in Organization Events API - supporting count(), eps() and epm() columns.

Required some manipulation on results to add the "()" back to the alias on the results from Snuba, because our SnQL parser doesnt allow them on aliases, but we want to convert back to the original column name after changing it under the hood.

Author: yuvmen

src/sentry/api/bases/organization_events.py

@@ -18,6 +18,7 @@
 from sentry.api.base import CURSOR_LINK_HEADER
 from sentry.api.bases import NoProjects
 from sentry.api.bases.organization import FilterParamsDateNotNull, OrganizationEndpoint
+from sentry.api.helpers.error_upsampling import are_all_projects_error_upsampled
 from sentry.api.helpers.mobile import get_readable_device_name
 from sentry.api.helpers.teams import get_teams
 from sentry.api.serializers.snuba import SnubaTSResultSerializer
@@ -340,6 +341,8 @@ def handle_results_with_meta(
             meta = results.get("meta", {})
             fields_meta = meta.get("fields", {})
 
+            self.handle_error_upsampling(project_ids, results)
+
             if standard_meta:
                 isMetricsData = meta.pop("isMetricsData", False)
                 isMetricsExtractedData = meta.pop("isMetricsExtractedData", False)
@@ -422,6 +425,37 @@ def handle_data(
 
         return results
 
+    def handle_error_upsampling(self, project_ids: Sequence[int], results: dict[str, Any]):
+        """
+        If the query is for error upsampled projects, we need to rename the fields to include the ()
+        and update the meta fields to reflect the new field names. This works around a limitation in
+        how aliases are handled in the SnQL parser.
+        """
+        if are_all_projects_error_upsampled(project_ids):
+            data = results.get("data", [])
+            fields_meta = results.get("meta", {}).get("fields", {})
+
+            for result in data:
+                if "count" in result:
+                    result["count()"] = result["count"]
+                    del result["count"]
+                if "eps" in result:
+                    result["eps()"] = result["eps"]
+                    del result["eps"]
+                if "epm" in result:
+                    result["epm()"] = result["epm"]
+                    del result["epm"]
+
+            if "count" in fields_meta:
+                fields_meta["count()"] = fields_meta["count"]
+                del fields_meta["count"]
+            if "eps" in fields_meta:
+                fields_meta["eps()"] = fields_meta["eps"]
+                del fields_meta["eps"]
+            if "epm" in fields_meta:
+                fields_meta["epm()"] = fields_meta["epm"]
+                del fields_meta["epm"]
+
     def handle_issues(
         self, results: Sequence[Any], project_ids: Sequence[int], organization: Organization
     ) -> None:

src/sentry/api/endpoints/organization_events.py

@@ -13,6 +13,10 @@
 from sentry.api.api_publish_status import ApiPublishStatus
 from sentry.api.base import region_silo_endpoint
 from sentry.api.bases import NoProjects, OrganizationEventsV2EndpointBase
+from sentry.api.helpers.error_upsampling import (
+    is_errors_query_for_error_upsampled_projects,
+    transform_query_columns_for_error_upsampling,
+)
 from sentry.api.paginator import GenericOffsetPaginator
 from sentry.api.utils import handle_query_errors
 from sentry.apidocs import constants as api_constants
@@ -315,9 +319,16 @@ def _data_fn(
             limit: int,
             query: str | None,
         ):
+            transform_alias_to_input_format = True
+            selected_columns = self.get_field_list(organization, request)
+            if is_errors_query_for_error_upsampled_projects(
+                snuba_params, organization, dataset, request
+            ):
+                selected_columns = transform_query_columns_for_error_upsampling(selected_columns)
+                transform_alias_to_input_format = False
             query_source = self.get_request_source(request)
             return dataset_query(
-                selected_columns=self.get_field_list(organization, request),
+                selected_columns=selected_columns,
                 query=query or "",
                 snuba_params=snuba_params,
                 equations=self.get_equation_list(organization, request),
@@ -329,7 +340,7 @@ def _data_fn(
                 auto_aggregations=True,
                 allow_metric_aggregates=allow_metric_aggregates,
                 use_aggregate_conditions=use_aggregate_conditions,
-                transform_alias_to_input_format=True,
+                transform_alias_to_input_format=transform_alias_to_input_format,
                 # Whether the flag is enabled or not, regardless of the referrer
                 has_metrics=use_metrics,
                 use_metrics_layer=batch_features.get("organizations:use-metrics-layer", False),

src/sentry/api/helpers/error_upsampling.py

@@ -44,21 +44,23 @@ def are_all_projects_error_upsampled(project_ids: Sequence[int]) -> bool:
     return result
 
 
-def transform_query_columns_for_error_upsampling(
-    query_columns: Sequence[str],
-) -> list[str]:
+def transform_query_columns_for_error_upsampling(query_columns: Sequence[str]) -> list[str]:
     """
     Transform aggregation functions to use sum(sample_weight) instead of count()
-    for error upsampling. Only called when all projects are allowlisted.
+    for error upsampling.
     """
     transformed_columns = []
     for column in query_columns:
         column_lower = column.lower().strip()
 
         if column_lower == "count()":
-            # Simple count becomes sum of sample weights
             transformed_columns.append("upsampled_count() as count")
 
+        elif column_lower == "eps()":
+            transformed_columns.append("upsampled_eps() as eps")
+
+        elif column_lower == "epm()":
+            transformed_columns.append("upsampled_epm() as epm")
         else:
             transformed_columns.append(column)
 

src/sentry/search/events/datasets/discover.py

@@ -1046,7 +1046,23 @@ def function_converter(self) -> Mapping[str, SnQLFunction]:
                         [Function("sum", [Function("ifNull", [Column("sample_weight"), 1])])],
                         alias,
                     ),
-                    default_result_type="number",
+                    default_result_type="integer",
+                ),
+                SnQLFunction(
+                    "upsampled_eps",
+                    snql_aggregate=lambda args, alias: function_aliases.resolve_upsampled_eps(
+                        args, alias, self.builder
+                    ),
+                    optional_args=[IntervalDefault("interval", 1, None)],
+                    default_result_type="rate",
+                ),
+                SnQLFunction(
+                    "upsampled_epm",
+                    snql_aggregate=lambda args, alias: function_aliases.resolve_upsampled_epm(
+                        args, alias, self.builder
+                    ),
+                    optional_args=[IntervalDefault("interval", 1, None)],
+                    default_result_type="rate",
                 ),
             ]
         }

src/sentry/search/events/datasets/function_aliases.py

@@ -379,6 +379,22 @@ def resolve_eps(
     return Function("divide", [Function("count", []), interval], alias)
 
 
+def resolve_upsampled_eps(
+    args: Mapping[str, str | Column | SelectType | int | float],
+    alias: str,
+    builder: BaseQueryBuilder,
+) -> SelectType:
+    if hasattr(builder, "interval"):
+        interval = builder.interval
+    else:
+        interval = args["interval"]
+    return Function(
+        "divide",
+        [Function("sum", [Function("ifNull", [Column("sample_weight"), 1])]), interval],
+        alias,
+    )
+
+
 def resolve_epm(
     args: Mapping[str, str | Column | SelectType | int | float],
     alias: str,
@@ -393,3 +409,22 @@ def resolve_epm(
         [Function("count", []), Function("divide", [interval, 60])],
         alias,
     )
+
+
+def resolve_upsampled_epm(
+    args: Mapping[str, str | Column | SelectType | int | float],
+    alias: str,
+    builder: BaseQueryBuilder,
+) -> SelectType:
+    if hasattr(builder, "interval"):
+        interval = builder.interval
+    else:
+        interval = args["interval"]
+    return Function(
+        "divide",
+        [
+            Function("sum", [Function("ifNull", [Column("sample_weight"), 1])]),
+            Function("divide", [interval, 60]),
+        ],
+        alias,
+    )