refactor handle_error_sampling

fix issues with double conversion and moved logic into helper

Author: yuvmen

src/sentry/api/bases/organization_events.py

@@ -18,7 +18,10 @@
 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_any_projects_error_upsampled
+from sentry.api.helpers.error_upsampling import (
+    are_any_projects_error_upsampled,
+    convert_fields_for_upsampling,
+)
 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
@@ -425,30 +428,14 @@ def handle_data(
 
     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 data and meta fields to reflect the new field names. This works around a limitation in
-        how aliases are handled in the SnQL parser.
+        If the query is for error upsampled projects, we convert various functions under the hood.
+        We need to rename these fields before returning the results to the client, to hide the conversion.
+        This is done here to work around a limitation in how aliases are handled in the SnQL parser.
         """
         if are_any_projects_error_upsampled(project_ids):
             data = results.get("data", [])
             fields_meta = results.get("meta", {}).get("fields", {})
-
-            function_conversions = {
-                "upsampled_count()": "count()",
-                "upsampled_eps()": "eps()",
-                "upsampled_epm()": "epm()",
-            }
-
-            # Go over each both data and meta, and convert function names to the non-upsampled version
-            for upsampled_function, count_function in function_conversions.items():
-                for result in data:
-                    if upsampled_function in result:
-                        result[count_function] = result[upsampled_function]
-                        del result[upsampled_function]
-
-                if upsampled_function in fields_meta:
-                    fields_meta[count_function] = fields_meta[upsampled_function]
-                    del fields_meta[upsampled_function]
+            convert_fields_for_upsampling(data, fields_meta)
 
     def handle_issues(
         self, results: Sequence[Any], project_ids: Sequence[int], organization: Organization

src/sentry/api/helpers/error_upsampling.py

@@ -10,6 +10,57 @@
 
 UPSAMPLED_ERROR_AGGREGATION = "upsampled_count"
 
+# Function key conversions for error upsampling results
+_FUNCTION_KEY_CONVERSIONS = {
+    "count()": "sample_count()",
+    "eps()": "sample_eps()",
+    "epm()": "sample_epm()",
+    "upsampled_count()": "count()",
+    "upsampled_eps()": "eps()",
+    "upsampled_epm()": "epm()",
+}
+
+# Pre-computed ordered keys to handle conversion conflicts
+# Keys that are targets of other conversions must be processed first
+_conversion_targets = set(_FUNCTION_KEY_CONVERSIONS.values())
+_ORDERED_CONVERSION_KEYS = sorted(
+    _FUNCTION_KEY_CONVERSIONS.keys(), key=lambda k: k not in _conversion_targets
+)
+
+
+def convert_fields_for_upsampling(data: list[dict[str, Any]], fields_meta: dict[str, str]) -> None:
+    """
+    Convert field names in query results for error upsampled projects.
+    This renames upsampled_* functions to their standard names and standard functions
+    to sample_* equivalents to hide the conversion from the client.
+
+    Args:
+        data: List of result dictionaries to modify in-place
+        fields_meta: Meta fields dictionary to modify in-place
+    """
+    # Collect keys that need conversion and exist in data
+    all_present_keys = set()
+    for result in data:
+        all_present_keys.update(result.keys())
+
+    # Filter the pre-ordered list to only include keys actually present
+    keys_to_convert = [key for key in _ORDERED_CONVERSION_KEYS if key in all_present_keys]
+
+    # Apply conversions to data
+    for result in data:
+        for original_key in keys_to_convert:
+            if original_key in result:
+                converted_key = _FUNCTION_KEY_CONVERSIONS[original_key]
+                result[converted_key] = result[original_key]
+                del result[original_key]
+
+    # Apply conversions to fields_meta
+    for original_key in keys_to_convert:
+        if original_key in fields_meta:
+            converted_key = _FUNCTION_KEY_CONVERSIONS[original_key]
+            fields_meta[converted_key] = fields_meta[original_key]
+            del fields_meta[original_key]
+
 
 def is_errors_query_for_error_upsampled_projects(
     snuba_params: SnubaParams,
@@ -55,6 +106,9 @@ def transform_query_columns_for_error_upsampling(
         "count()": "upsampled_count()",
         "eps()": "upsampled_eps()",
         "epm()": "upsampled_epm()",
+        "sample_count()": "count()",
+        "sample_eps()": "eps()",
+        "sample_epm()": "epm()",
     }
 
     transformed_columns = []