From e8fc76e74bbd181d2bfb9122255cea48c30c75ca Mon Sep 17 00:00:00 2001
From: Riya Chakraborty <riya.chakraborty@sentry.io>
Date: Mon, 13 May 2024 15:45:11 -0700
Subject: [PATCH] add no percentiles

---
 src/sentry/options/defaults.py                         |  8 ++++++++
 .../sentry_metrics/aggregation_option_registry.py      | 10 +++++++---
 src/sentry/sentry_metrics/consumers/indexer/batch.py   |  4 +++-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py
index 8619a4fe0116f6..b35a4d81afb56c 100644
--- a/src/sentry/options/defaults.py
+++ b/src/sentry/options/defaults.py
@@ -1237,6 +1237,14 @@
     "sentry-metrics.indexer.reconstruct.enable-orjson", default=0.0, flags=FLAG_AUTOMATOR_MODIFIABLE
 )
 
+# Option to remove support for percentiles on a per-org basis.
+# Add the org_id to list to disable percentiles.
+register(
+    "sentry-metrics.drop-percentiles.per-org",
+    default=[],
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 # Global and per-organization limits on the writes to the string indexer's DB.
 #
 # Format is a list of dictionaries of format {
diff --git a/src/sentry/sentry_metrics/aggregation_option_registry.py b/src/sentry/sentry_metrics/aggregation_option_registry.py
index 0a03211930f5c4..571a42c5d1b7b4 100644
--- a/src/sentry/sentry_metrics/aggregation_option_registry.py
+++ b/src/sentry/sentry_metrics/aggregation_option_registry.py
@@ -8,6 +8,7 @@
 class AggregationOption(Enum):
     HIST = "hist"
     TEN_SECOND = "ten_second"
+    NO_PERCENTILE = "no_percentile"
 
 
 class TimeWindow(Enum):
@@ -25,11 +26,14 @@ class TimeWindow(Enum):
 USE_CASE_AGG_OPTION = {UseCaseID.CUSTOM: {AggregationOption.TEN_SECOND: TimeWindow.SEVEN_DAYS}}
 
 
-def get_aggregation_options(mri: str) -> dict[AggregationOption, TimeWindow] | None:
+def get_aggregation_options(mri: str, org_id: int) -> dict[AggregationOption, TimeWindow] | None:
     use_case_id: UseCaseID = extract_use_case_id(mri)
 
-    # We check first if the particular metric ID has a specified aggregation
-    if mri in METRIC_ID_AGG_OPTION:
+    # We check first if the org ID has disabled percentiles
+    if org_id in options.get("sentry-metrics.drop-percentiles.per-org"):
+        return {AggregationOption.NO_PERCENTILE: TimeWindow.NINETY_DAYS}
+    # We then first if the particular metric ID has a specified aggregation
+    elif mri in METRIC_ID_AGG_OPTION:
         return METRIC_ID_AGG_OPTION.get(mri)
     # And move to the use case if not
     elif options.get("sentry-metrics.10s-granularity") and (use_case_id in USE_CASE_AGG_OPTION):
diff --git a/src/sentry/sentry_metrics/consumers/indexer/batch.py b/src/sentry/sentry_metrics/consumers/indexer/batch.py
index 7dcb54b97300c3..710fc13bf928b6 100644
--- a/src/sentry/sentry_metrics/consumers/indexer/batch.py
+++ b/src/sentry/sentry_metrics/consumers/indexer/batch.py
@@ -488,7 +488,9 @@ def reconstruct_messages(
                         "value": old_payload_value["value"],
                         "sentry_received_timestamp": sentry_received_timestamp,
                     }
-                    if aggregation_options := get_aggregation_options(old_payload_value["name"]):
+                    if aggregation_options := get_aggregation_options(
+                        old_payload_value["name"], old_payload_value["org_id"]
+                    ):
                         # TODO: This should eventually handle multiple aggregation options
                         option = list(aggregation_options.items())[0][0]
                         assert option is not None