fix: Fix some typing problems around dynamic sampling (#51256)

Co-authored-by: getsantry[bot] <66042841+getsantry[bot]@users.noreply.github.com>

pyproject.toml

@@ -263,7 +263,6 @@ module = [
     "sentry.api.endpoints.organization_events_meta",
     "sentry.api.endpoints.organization_events_span_ops",
     "sentry.api.endpoints.organization_events_spans_performance",
-    "sentry.api.endpoints.organization_events_starfish",
     "sentry.api.endpoints.organization_events_stats",
     "sentry.api.endpoints.organization_events_trace",
     "sentry.api.endpoints.organization_events_trends",
@@ -863,7 +862,6 @@ module = [
     "sentry.plugins.validators.url",
     "sentry.post_process_forwarder.post_process_forwarder",
     "sentry.profiles.consumers",
-    "sentry.profiles.flamegraph",
     "sentry.profiles.task",
     "sentry.profiles.utils",
     "sentry.projectoptions.defaults",
@@ -1297,7 +1295,6 @@ module = [
     "tests.sentry.api.endpoints.test_project_artifact_bundle_file_details",
     "tests.sentry.api.endpoints.test_project_artifact_lookup",
     "tests.sentry.api.endpoints.test_project_details",
-    "tests.sentry.api.endpoints.test_project_dynamic_sampling",
     "tests.sentry.api.endpoints.test_project_group_stats",
     "tests.sentry.api.endpoints.test_project_ownership",
     "tests.sentry.api.endpoints.test_project_release_details",

src/sentry/api/endpoints/organization_profiling_functions.py

@@ -108,7 +108,7 @@ class OrganizationProfilingFunctionTrendsEndpoint(OrganizationEventsV2EndpointBa
                 params=params,
                 orderby=["-count()"],
                 limit=TOP_FUNCTIONS_LIMIT,
-                referrer=Referrer.API_PROFILING_FUNCTION_TRENDS_TOP_EVENTS.value,  # type: ignore[attr-defined]
+                referrer=Referrer.API_PROFILING_FUNCTION_TRENDS_TOP_EVENTS.value,
                 use_aggregate_conditions=True,
                 transform_alias_to_input_format=True,
             )
@@ -126,7 +126,7 @@ class OrganizationProfilingFunctionTrendsEndpoint(OrganizationEventsV2EndpointBa
                 top_events=top_functions,
                 organization=organization,
                 zerofill_results=zerofill_results,
-                referrer=Referrer.API_PROFILING_FUNCTION_TRENDS_STATS.value,  # type: ignore[attr-defined]
+                referrer=Referrer.API_PROFILING_FUNCTION_TRENDS_STATS.value,
                 # this ensures the result key is formatted as `{project.id},{fingerprint}`
                 # in order to be compatible with the trends service
                 result_key_order=["project.id", "fingerprint"],

src/sentry/dynamic_sampling/tasks/boost_low_volume_projects.py

@@ -182,7 +182,7 @@ def fetch_projects_with_total_root_transaction_count_and_rates(
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value,
         )["data"]
         count = len(data)
         more_results = count > CHUNK_SIZE
@@ -228,7 +228,7 @@ def adjust_sample_rates_of_projects(
             org_id, None, "boost_low_volume_projects", "sliding_window_org", sample_rate
         )
     else:
-        sample_rate = quotas.get_blended_sample_rate(organization_id=org_id)  # type:ignore
+        sample_rate = quotas.backend.get_blended_sample_rate(organization_id=org_id)
         log_sample_rate_source(
             org_id, None, "boost_low_volume_projects", "blended_sample_rate", sample_rate
         )

src/sentry/dynamic_sampling/tasks/boost_low_volume_transactions.py

@@ -142,7 +142,7 @@ def boost_low_volume_transactions_of_project(project_transactions: ProjectTransa
         organization = None
 
     # By default, this bias uses the blended sample rate.
-    sample_rate = quotas.get_blended_sample_rate(organization_id=org_id)  # type:ignore
+    sample_rate = quotas.backend.get_blended_sample_rate(organization_id=org_id)
 
     # In case we have specific feature flags enabled, we will change the sample rate either basing ourselves
     # on sliding window per project or per org.
@@ -266,7 +266,7 @@ def get_orgs_with_project_counts(
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,
         )["data"]
         count = len(data)
         more_results = count > CHUNK_SIZE
@@ -347,7 +347,7 @@ def fetch_project_transaction_totals(org_ids: List[int]) -> Iterator[ProjectTran
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_TRANSACTION_TOTALS.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_TRANSACTION_TOTALS.value,
         )["data"]
         count = len(data)
         more_results = count > CHUNK_SIZE
@@ -450,7 +450,7 @@ def fetch_transactions_with_total_volumes(
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,
         )["data"]
         count = len(data)
         more_results = count > CHUNK_SIZE

src/sentry/dynamic_sampling/tasks/common.py

@@ -79,7 +79,7 @@ def get_active_orgs_with_projects_counts(
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_FETCH_PROJECTS_WITH_COUNT_PER_TRANSACTION.value,
         )["data"]
         count = len(data)
         more_results = count > CHUNK_SIZE

src/sentry/dynamic_sampling/tasks/recalibrate_orgs.py

@@ -160,7 +160,7 @@ def get_active_orgs(max_orgs: int, time_interval: timedelta) -> Generator[List[i
         )
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_GET_ACTIVE_ORGS.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_GET_ACTIVE_ORGS.value,
         )["data"]
         count = len(data)
         more_results = count > max_orgs
@@ -229,7 +229,7 @@ def fetch_org_volumes(
     )
     data = raw_snql_query(
         request,
-        referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_GET_ORG_TRANSACTION_VOLUMES.value,  # type:ignore
+        referrer=Referrer.DYNAMIC_SAMPLING_COUNTERS_GET_ORG_TRANSACTION_VOLUMES.value,
     )["data"]
 
     for row in data:

src/sentry/dynamic_sampling/tasks/sliding_window.py

@@ -188,7 +188,7 @@ def fetch_projects_with_total_root_transactions_count(
 
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECTS_WITH_COUNT_PER_ROOT.value,
         )["data"]
 
         count = len(data)

src/sentry/dynamic_sampling/tasks/sliding_window_org.py

@@ -130,7 +130,7 @@ def fetch_orgs_with_total_root_transactions_count(
 
         data = raw_snql_query(
             request,
-            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_ORGS_WITH_COUNT_PER_ROOT.value,  # type:ignore
+            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_ORGS_WITH_COUNT_PER_ROOT.value,
         )["data"]
 
         count = len(data)

src/sentry/snuba/referrer.py

@@ -2,7 +2,7 @@ import enum
 import logging
 from enum import Enum, unique
 from itertools import chain
-from typing import Mapping, Optional
+from typing import TYPE_CHECKING, Mapping, Optional
 
 from sentry.tsdb.base import TSDBModel
 from sentry.utils import metrics
@@ -619,10 +619,21 @@ TSDBModelSuffixReferrer = enum.Enum(
 )
 
 
-Referrer = enum.Enum(
-    "Referrer",
-    [(i.name, i.value) for i in chain(ReferrerBase, TSDBModelReferrer, TSDBModelSuffixReferrer)],
-)
+if TYPE_CHECKING:
+    # This type is hot garbage (enums don't support subclassing) and not
+    # suitable for runtime, but it's good enough for mypy so that it can find
+    # properties on the Referrer object
+    class Referrer(ReferrerBase, TSDBModelReferrer, TSDBModelSuffixReferrer):
+        pass
+
+else:
+    Referrer = enum.Enum(
+        "Referrer",
+        [
+            (i.name, i.value)
+            for i in chain(ReferrerBase, TSDBModelReferrer, TSDBModelSuffixReferrer)
+        ],
+    )
 
 
 def validate_referrer(referrer: Optional[str]):

tests/sentry/dynamic_sampling/tasks/test_tasks.py

@@ -124,7 +124,7 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
     def forecasted_volume_side_effect(*args, **kwargs):
         return kwargs["volume"]
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_projects_simple(
         self,
         get_blended_sample_rate,
@@ -159,7 +159,7 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
         }
         assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0}
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_projects_simple_with_empty_project(
         self,
         get_blended_sample_rate,
@@ -196,8 +196,8 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
         assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0}
         assert generate_rules(proj_e)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0}
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_boost_low_volume_projects_simple_with_sliding_window_org_from_cache(
         self,
@@ -238,8 +238,8 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
         }
         assert generate_rules(proj_d)[0]["samplingValue"] == {"type": "sampleRate", "value": 1.0}
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_boost_low_volume_projects_simple_with_sliding_window_org_from_sync(
         self,
@@ -284,8 +284,8 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
     @patch(
         "sentry.dynamic_sampling.tasks.boost_low_volume_projects.schedule_invalidate_project_config"
     )
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_config_invalidation_when_sample_rates_change(
         self,
@@ -316,8 +316,8 @@ class TestBoostLowVolumeProjectsTasks(TasksTestCase):
     @patch(
         "sentry.dynamic_sampling.tasks.boost_low_volume_projects.schedule_invalidate_project_config"
     )
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_config_invalidation_when_sample_rates_do_not_change(
         self,
@@ -454,7 +454,7 @@ class TestBoostLowVolumeTransactionsTasks(TasksTestCase):
             )
         )
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_transactions_with_blended_sample_rate(self, get_blended_sample_rate):
         """
         Create orgs projects & transactions and then check that the task creates rebalancing data
@@ -479,7 +479,7 @@ class TestBoostLowVolumeTransactionsTasks(TasksTestCase):
                     )  # check we have some rate calculated for each transaction
                 assert global_rate == BLENDED_RATE
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_transactions_with_sliding_window(self, get_blended_sample_rate):
         """
         Create orgs projects & transactions and then check that the task creates rebalancing data
@@ -511,7 +511,7 @@ class TestBoostLowVolumeTransactionsTasks(TasksTestCase):
                         )  # check we have some rate calculated for each transaction
                     assert global_rate == used_sample_rate
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_transactions_with_sliding_window_org(self, get_blended_sample_rate):
         """
         Create orgs projects & transactions and then check that the task creates rebalancing data
@@ -560,7 +560,7 @@ class TestBoostLowVolumeTransactionsTasks(TasksTestCase):
 
                         assert global_rate == used_sample_rate
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_boost_low_volume_transactions_partial(self, get_blended_sample_rate):
         """
         Test the V2 algorithm is used, only specified projects are balanced and the
@@ -650,7 +650,7 @@ class TestRecalibrateOrgsTasks(TasksTestCase):
                     org_id=org.id,
                 )
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_rebalance_orgs(self, get_blended_sample_rate):
         """
         Test that the org are going to be rebalanced at 20%
@@ -701,7 +701,7 @@ class TestRecalibrateOrgsTasks(TasksTestCase):
                 # half it again to 0.25
                 assert float(val) == 0.25
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_rebalance_rules(self, get_blended_sample_rate):
         """
         Test that we pass rebalancing values all the way to the rules
@@ -789,8 +789,8 @@ class TestSlidingWindowTasks(TasksTestCase):
             is not None
         )
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_sliding_window_with_multiple_projects(
         self,
@@ -845,8 +845,8 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 1.0,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_sliding_window_with_none_sampling_tier(
         self,
@@ -879,7 +879,7 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 0.5,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_sliding_window_with_forecasting_error(
         self, extrapolate_monthly_volume, get_blended_sample_rate
@@ -902,7 +902,7 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 0.9,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     @patch("sentry.dynamic_sampling.tasks.common.compute_sliding_window_sample_rate")
     def test_sliding_window_with_sample_rate_computation_error(
         self, compute_sliding_window_sample_rate, get_blended_sample_rate
@@ -925,8 +925,8 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 0.9,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_sliding_window_with_project_without_metrics(
         self,
@@ -958,8 +958,8 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 1.0,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_sliding_window_with_all_projects_without_metrics(
         self,
@@ -989,7 +989,7 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 1.0,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_sliding_window_with_none_window_size(self, get_blended_sample_rate):
         get_blended_sample_rate.return_value = 0.5
 
@@ -1010,8 +1010,8 @@ class TestSlidingWindowTasks(TasksTestCase):
                 }
 
     @patch("sentry.dynamic_sampling.tasks.sliding_window.schedule_invalidate_project_config")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_config_invalidation_when_sample_rates_change(
         self,
@@ -1043,8 +1043,8 @@ class TestSlidingWindowTasks(TasksTestCase):
         assert schedule_invalidate_project_config.call_count == 2
 
     @patch("sentry.dynamic_sampling.tasks.sliding_window.schedule_invalidate_project_config")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_config_invalidation_when_sample_rates_do_not_change(
         self,
@@ -1076,8 +1076,8 @@ class TestSlidingWindowTasks(TasksTestCase):
         schedule_invalidate_project_config.assert_not_called()
 
     @patch("sentry.dynamic_sampling.tasks.sliding_window.schedule_invalidate_project_config")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_transaction_sampling_tier_for_volume")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_transaction_sampling_tier_for_volume")
     @patch("sentry.dynamic_sampling.tasks.common.extrapolate_monthly_volume")
     def test_cache_deletion_when_project_has_no_more_metrics(
         self,
@@ -1127,7 +1127,7 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 1.0,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_sliding_window_with_new_project(self, get_blended_sample_rate):
         get_blended_sample_rate.return_value = 0.5
 
@@ -1147,7 +1147,7 @@ class TestSlidingWindowTasks(TasksTestCase):
                 "value": 1.0,
             }
 
-    @patch("sentry.dynamic_sampling.rules.base.quotas.get_blended_sample_rate")
+    @patch("sentry.quotas.backend.get_blended_sample_rate")
     def test_sliding_window_with_new_org(self, get_blended_sample_rate):
         get_blended_sample_rate.return_value = 0.5