feat(CapMan): More tsdb and legacy queries with tenant_ids (#45405)

rahul-kumar-saini · web-flow · commit 0183b0753414 · 2023-03-07T10:41:07.000-08:00
### Overview - Added `organization_id` to more Snuba Requests, specifically targeting `tsdb-modelid:300` and `tsdb-modelid:100` - Legacy queries now pass `tenant_ids` to Snuba as well - Removed hardcoded referrer in tenant_ids, moved to populate referrer in query functions - More context: #44788
diff --git a/src/sentry/api/endpoints/group_hashes_split.py b/src/sentry/api/endpoints/group_hashes_split.py
@@ -159,14 +159,14 @@ def _get_full_hierarchical_hashes(group: Group, hash: str) -> Optional[Sequence[
             ]
         )
     )
-    referrer = "group_split.get_full_hierarchical_hashes"
+
     request = SnubaRequest(
         dataset="events",
         app_id="grouping",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": group.project.organization_id},
+        tenant_ids={"organization_id": group.project.organization_id},
     )
-    data = snuba.raw_snql_query(request, referrer)["data"]
+    data = snuba.raw_snql_query(request, "group_split.get_full_hierarchical_hashes")["data"]
     if not data:
         return None
 
@@ -390,17 +390,13 @@ def _render_trees(group: Group, user):
     )
 
     rv = []
-    referrer = "api.group_split.render_grouping_tree"
     request = SnubaRequest(
         dataset="events",
         app_id="grouping",
         query=query,
-        tenant_ids={
-            "referrer": referrer,
-            "organization_id": group.project.organization_id,
-        },
+        tenant_ids={"organization_id": group.project.organization_id},
     )
-    for row in snuba.raw_snql_query(request, referrer)["data"]:
+    for row in snuba.raw_snql_query(request, "api.group_split.render_grouping_tree")["data"]:
         if len(row["hash_slice"]) == 0:
             hash = row["primary_hash"]
             parent_hash = child_hash = None
diff --git a/src/sentry/api/endpoints/grouping_level_new_issues.py b/src/sentry/api/endpoints/grouping_level_new_issues.py
@@ -195,14 +195,13 @@ def _query_snuba(group: Group, id: int, offset=None, limit=None):
     if limit is not None:
         query = query.set_limit(limit)
 
-    referrer = "api.group_hashes_levels.get_level_new_issues"
     request = SnubaRequest(
         dataset="events",
         app_id="grouping",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": group.project.organization_id},
+        tenant_ids={"organization_id": group.project.organization_id},
     )
-    return snuba.raw_snql_query(request, referrer)["data"]
+    return snuba.raw_snql_query(request, "api.group_hashes_levels.get_level_new_issues")["data"]
 
 
 def _process_snuba_results(query_res, group: Group, id: int, user):
diff --git a/src/sentry/api/endpoints/grouping_levels.py b/src/sentry/api/endpoints/grouping_levels.py
@@ -127,14 +127,13 @@ def get_levels_overview(group: Group):
         .set_where(_get_group_filters(group))
         .set_groupby([Column("primary_hash")])
     )
-    referrer = "api.group_hashes_levels.get_levels_overview"
     request = SnubaRequest(
         dataset="events",
         app_id="grouping",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": group.project.organization_id},
+        tenant_ids={"organization_id": group.project.organization_id},
     )
-    res = snuba.raw_snql_query(request, referrer)
+    res = snuba.raw_snql_query(request, "api.group_hashes_levels.get_levels_overview")
 
     if not res["data"]:
         raise NoEvents()
diff --git a/src/sentry/api/endpoints/organization_events_span_ops.py b/src/sentry/api/endpoints/organization_events_span_ops.py
@@ -38,10 +38,9 @@ def data_fn(offset: int, limit: int) -> Any:
                 offset=offset,
                 orderby="-count",
             )
-            referrer = "api.organization-events-span-ops"
             snql_query = builder.get_snql_query()
-            snql_query.tenant_ids = {"referrer": referrer, "organization_id": organization.id}
-            results = raw_snql_query(snql_query, referrer)
+            snql_query.tenant_ids = {"organization_id": organization.id}
+            results = raw_snql_query(snql_query, "api.organization-events-span-ops")
             return [SpanOp(op=row["spans_op"], count=row["count"]) for row in results["data"]]
 
         with self.handle_query_errors():
diff --git a/src/sentry/api/endpoints/project_dynamic_sampling.py b/src/sentry/api/endpoints/project_dynamic_sampling.py
@@ -110,16 +110,15 @@ def __project_stats_query(self, query, projects_in_org, org_id, query_time_range
                 alias="root_count",
             )
         ]
-        referrer = Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECT_STATS.value
         snuba_query = snuba_query.set_select(snuba_query.select + extra_select)
         data = raw_snql_query(
             SnubaRequest(
                 dataset=Dataset.Discover.value,
                 app_id="default",
                 query=snuba_query,
-                tenant_ids={"referrer": referrer, "organization_id": org_id},
+                tenant_ids={"organization_id": org_id},
             ),
-            referrer,
+            Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECT_STATS.value,
         )
         return builder.process_results(data)["data"]
 
@@ -266,8 +265,6 @@ def __fetch_randomly_sampled_transactions(self, project, query, sample_size, que
             ]
         )
 
-        referrer = Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_TRANSACTIONS.value
-
         snuba_query = snuba_query.set_groupby(
             snuba_query.groupby + [Column("modulo_num"), Column("contexts.key")]
         )
@@ -277,9 +274,9 @@ def __fetch_randomly_sampled_transactions(self, project, query, sample_size, que
                 dataset=Dataset.Discover.value,
                 app_id="default",
                 query=snuba_query,
-                tenant_ids={"referrer": referrer, "organization_id": project.organization_id},
+                tenant_ids={"organization_id": project.organization_id},
             ),
-            referrer,
+            Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_TRANSACTIONS.value,
         )["data"]
         return data
 
diff --git a/src/sentry/api/serializers/models/group.py b/src/sentry/api/serializers/models/group.py
@@ -1007,6 +1007,10 @@ def _execute_error_seen_stats_query(
         filters = {"project_id": project_ids, "group_id": group_ids}
         if environment_ids:
             filters["environment"] = environment_ids
+
+        org_id = item_list[0].project.organization_id if item_list else None
+        tenant_ids = {"organization_id": org_id} if org_id else dict()
+
         return aliased_query(
             dataset=Dataset.Events,
             start=start,
@@ -1016,6 +1020,7 @@ def _execute_error_seen_stats_query(
             filter_keys=filters,
             aggregations=aggregations,
             referrer="serializers.GroupSerializerSnuba._execute_error_seen_stats_query",
+            tenant_ids=tenant_ids,
         )
 
     @staticmethod
diff --git a/src/sentry/digests/notifications.py b/src/sentry/digests/notifications.py
@@ -76,6 +76,7 @@ def fetch_state(project: Project, records: Sequence[Record]) -> Mapping[str, Any
     end = records[0].datetime
 
     groups = Group.objects.in_bulk(record.value.event.group_id for record in records)
+    tenant_ids = {"organization_id": project.organization_id}
     return {
         "project": project,
         "groups": groups,
@@ -87,10 +88,14 @@ def fetch_state(project: Project, records: Sequence[Record]) -> Mapping[str, Any
             list(groups.keys()),
             start,
             end,
-            tenant_ids={"organization_id": project.organization_id},
+            tenant_ids=tenant_ids,
         ),
         "user_counts": tsdb.get_distinct_counts_totals(
-            tsdb.models.users_affected_by_group, list(groups.keys()), start, end
+            tsdb.models.users_affected_by_group,
+            list(groups.keys()),
+            start,
+            end,
+            tenant_ids=tenant_ids,
         ),
     }
 
diff --git a/src/sentry/ingest/transaction_clusterer/datasource/snuba.py b/src/sentry/ingest/transaction_clusterer/datasource/snuba.py
@@ -12,7 +12,6 @@ def fetch_unique_transaction_names(
     project: Project, time_range: Tuple[datetime, datetime], limit: int
 ) -> Iterable[str]:
     then, now = time_range
-    referrer = "src.sentry.ingest.transaction_clusterer"
     snuba_request = Request(
         "transactions",
         app_id="transactions",
@@ -28,8 +27,8 @@ def fetch_unique_transaction_names(
             groupby=[Column("transaction")],
             limit=Limit(limit),
         ),
-        tenant_ids={"referrer": referrer, "organization_id": project.organization_id},
+        tenant_ids={"organization_id": project.organization_id},
     )
-    snuba_response = raw_snql_query(snuba_request, referrer)
+    snuba_response = raw_snql_query(snuba_request, "src.sentry.ingest.transaction_clusterer")
 
     return (row["transaction"] for row in snuba_response["data"])
diff --git a/src/sentry/models/groupsnooze.py b/src/sentry/models/groupsnooze.py
@@ -119,6 +119,7 @@ def test_user_rates(self):
             keys=[self.group_id],
             start=start,
             end=end,
+            tenant_ids={"organization_id": self.group.project.organization_id},
         )[self.group_id]
 
         if rate >= self.user_count:
diff --git a/src/sentry/release_health/release_monitor/metrics.py b/src/sentry/release_health/release_monitor/metrics.py
@@ -147,15 +147,16 @@ def fetch_project_release_health_totals(
                     .set_limit(self.CHUNK_SIZE + 1)
                     .set_offset(offset)
                 )
-                referrer = "release_monitor.fetch_project_release_health_totals"
                 request = Request(
                     dataset=Dataset.Metrics.value,
                     app_id="release_health",
                     query=query,
-                    tenant_ids={"referrer": referrer, "organization_id": org_id},
+                    tenant_ids={"organization_id": org_id},
                 )
                 with metrics.timer("release_monitor.fetch_project_release_health_totals.query"):
-                    data = raw_snql_query(request, referrer)["data"]
+                    data = raw_snql_query(
+                        request, "release_monitor.fetch_project_release_health_totals"
+                    )["data"]
                     count = len(data)
                     more_results = count > self.CHUNK_SIZE
                     offset += self.CHUNK_SIZE
diff --git a/src/sentry/rules/conditions/event_frequency.py b/src/sentry/rules/conditions/event_frequency.py
@@ -258,6 +258,7 @@ def query_hook(
             environment_id=environment_id,
             use_cache=True,
             jitter_value=event.group_id,
+            tenant_ids={"organization_id": event.group.project.organization_id},
         )
         return totals[event.group_id]
 
diff --git a/src/sentry/snuba/metrics/fields/base.py b/src/sentry/snuba/metrics/fields/base.py
@@ -149,7 +149,7 @@ def run_metrics_query(
         dataset=Dataset.Metrics.value,
         app_id="metrics",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": org_id},
+        tenant_ids={"organization_id": org_id},
     )
     result = raw_snql_query(request, referrer, use_cache=True)
     return cast(List[SnubaDataType], result["data"])
diff --git a/src/sentry/snuba/sessions.py b/src/sentry/snuba/sessions.py
@@ -155,14 +155,13 @@ def _check_releases_have_health_data(
             Condition(Column("release"), Op.IN, release_versions),
         ],
     )
-    referrer = "snuba.sessions.check_releases_have_health_data"
     request = Request(
         dataset="sessions",
         app_id="default",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": organization_id},
+        tenant_ids={"organization_id": organization_id},
     )
-    data = snuba.raw_snql_query(request, referrer)["data"]
+    data = snuba.raw_snql_query(request, "snuba.sessions.check_releases_have_health_data")["data"]
     return {row["release"] for row in data}
 
 
@@ -269,14 +268,13 @@ def _get_project_releases_count(
         where=where,
         having=having,
     )
-    referrer = "snuba.sessions.get_project_releases_count"
     request = Request(
         dataset="sessions",
         app_id="default",
         query=query,
-        tenant_ids={"referrer": referrer, "organization_id": organization_id},
+        tenant_ids={"organization_id": organization_id},
     )
-    data = snuba.raw_snql_query(request, referrer)["data"]
+    data = snuba.raw_snql_query(request, "snuba.sessions.get_project_releases_count")["data"]
     return data[0]["count"] if data else 0
 
 
diff --git a/src/sentry/tagstore/snuba/backend.py b/src/sentry/tagstore/snuba/backend.py
@@ -745,6 +745,7 @@ def get_release_tags(self, organization_id, project_ids, environment_id, version
             aggregations=aggregations,
             orderby="-times_seen",
             referrer="tagstore.get_release_tags",
+            tenant_ids={"organization_id": organization_id},
         )
 
         values = []
diff --git a/src/sentry/tsdb/base.py b/src/sentry/tsdb/base.py
@@ -461,6 +461,7 @@ def get_distinct_counts_totals(
         environment_id=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         """
         Count distinct items during a time range.
diff --git a/src/sentry/tsdb/dummy.py b/src/sentry/tsdb/dummy.py
@@ -57,6 +57,7 @@ def get_distinct_counts_totals(
         environment_id=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         self.validate_arguments([model], [environment_id])
         return {k: 0 for k in keys}
diff --git a/src/sentry/tsdb/inmemory.py b/src/sentry/tsdb/inmemory.py
@@ -138,6 +138,7 @@ def get_distinct_counts_totals(
         environment_id=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         self.validate_arguments([model], [environment_id])
 
diff --git a/src/sentry/tsdb/redis.py b/src/sentry/tsdb/redis.py
@@ -455,6 +455,7 @@ def get_distinct_counts_totals(
         environment_id=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         """
         Count distinct items during a time range.
diff --git a/src/sentry/tsdb/snuba.py b/src/sentry/tsdb/snuba.py
@@ -326,6 +326,7 @@ def get_data(
                 conditions,
                 use_cache,
                 jitter_value,
+                tenant_ids,
             )
 
     def __get_data_snql(
@@ -442,8 +443,6 @@ def __get_data_snql(
                     ]
 
             referrer = f"tsdb-modelid:{model.value}"
-            tenant_ids = tenant_ids or dict()
-            tenant_ids["referrer"] = referrer
             snql_request = Request(
                 dataset=model_dataset.value,
                 app_id="tsdb.get_data",
@@ -456,7 +455,7 @@ def __get_data_snql(
                     granularity=Granularity(rollup),
                     limit=Limit(limit),
                 ),
-                tenant_ids=tenant_ids,
+                tenant_ids=tenant_ids or dict(),
             )
             query_result = raw_snql_query(snql_request, referrer=referrer, use_cache=use_cache)
             if manual_group_on_time:
@@ -496,6 +495,7 @@ def __get_data_legacy(
         conditions=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         """
         Normalizes all the TSDB parameters and sends a query to snuba.
@@ -602,6 +602,7 @@ def __get_data_legacy(
                 referrer=f"tsdb-modelid:{model.value}",
                 is_grouprelease=(model == TSDBModel.frequent_releases_by_group),
                 use_cache=use_cache,
+                tenant_ids=tenant_ids or dict(),
             )
             if model_query_settings.selected_columns:
                 result = query_func_without_selected_columns(
@@ -771,6 +772,7 @@ def get_distinct_counts_totals(
         environment_id=None,
         use_cache=False,
         jitter_value=None,
+        tenant_ids=None,
     ):
         return self.get_data(
             model,
@@ -782,6 +784,7 @@ def get_distinct_counts_totals(
             aggregation="uniq",
             use_cache=use_cache,
             jitter_value=jitter_value,
+            tenant_ids=tenant_ids,
         )
 
     def get_distinct_counts_union(
diff --git a/src/sentry/utils/snuba.py b/src/sentry/utils/snuba.py
