Просмотр исходного кода

feat(CapMan): Pass `tenant_ids` to Snuba (#44788)

### Overview
- As part of the [Capacity
Management](https://www.notion.so/sentry/Clickhouse-Capacity-Management-df585f5b7ae4453a9e635e1520b4a211)
project over on
[S<sup>3</sup>](https://www.notion.so/sentry/Search-Storage-Streaming-a5d224329fe148e4b27ff228339258b1),
we're working on getting more information out of per Snuba Request in
order to properly categorize and eventually set precise limits on API
usage by a number of different variables (we're calling them tenants)

### Changes
- The first two tenant IDs we're implementing are `organization_id` and
`referrer` (referrer already existed but this PR adds it to the
dictionary)
- This PR updates the Snuba SDK version in Sentry and starts the process
of consolidating these tenants into a single dictionary called
`tenant_ids` per request
- I've extracted `organization_id` from some easy use cases of the Snuba
API and will be working on more complicated ones in future PRs

### Notes
- Eventually, we will need Sentry's backend teams to help figure out how
to get more of these tenant IDs into the requests if/when things get too
complicated
Rahul Kumar Saini 2 лет назад
Родитель
Сommit
d7a91df4ed

+ 1 - 1
requirements-base.txt

@@ -57,7 +57,7 @@ rfc3986-validator>=0.1.1
 sentry-arroyo>=2.6.0
 sentry-relay>=0.8.18
 sentry-sdk>=1.15.0
-snuba-sdk>=1.0.3
+snuba-sdk>=1.0.5
 simplejson>=3.17.6
 statsd>=3.3
 structlog>=22

+ 1 - 1
requirements-dev-frozen.txt

@@ -160,7 +160,7 @@ sentry-sdk==1.15.0
 simplejson==3.17.6
 six==1.16.0
 sniffio==1.2.0
-snuba-sdk==1.0.3
+snuba-sdk==1.0.5
 sortedcontainers==2.4.0
 soupsieve==2.3.2.post1
 sqlparse==0.3.0

+ 1 - 1
requirements-frozen.txt

@@ -113,7 +113,7 @@ sentry-sdk==1.15.0
 simplejson==3.17.6
 six==1.16.0
 sniffio==1.2.0
-snuba-sdk==1.0.3
+snuba-sdk==1.0.5
 sortedcontainers==2.4.0
 soupsieve==2.3.2.post1
 sqlparse==0.3.0

+ 20 - 8
src/sentry/api/endpoints/group_hashes_split.py

@@ -5,6 +5,7 @@ import sentry_sdk
 from django.db import transaction
 from rest_framework.request import Request
 from rest_framework.response import Response
+from snuba_sdk import Request as SnubaRequest
 from snuba_sdk.conditions import Condition, Op
 from snuba_sdk.orderby import Direction, OrderBy
 from snuba_sdk.query import Column, Entity, Function, Query
@@ -158,10 +159,14 @@ def _get_full_hierarchical_hashes(group: Group, hash: str) -> Optional[Sequence[
             ]
         )
     )
-    request = Request(dataset="events", app_id="grouping", query=query)
-    data = snuba.raw_snql_query(request, referrer="group_split.get_full_hierarchical_hashes")[
-        "data"
-    ]
+    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},
+    )
+    data = snuba.raw_snql_query(request, referrer)["data"]
     if not data:
         return None
 
@@ -385,10 +390,17 @@ def _render_trees(group: Group, user):
     )
 
     rv = []
-    request = Request(dataset="events", app_id="grouping", query=query)
-    for row in snuba.raw_snql_query(request, referrer="api.group_split.render_grouping_tree")[
-        "data"
-    ]:
+    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,
+        },
+    )
+    for row in snuba.raw_snql_query(request, referrer)["data"]:
         if len(row["hash_slice"]) == 0:
             hash = row["primary_hash"]
             parent_hash = child_hash = None

+ 17 - 5
src/sentry/api/endpoints/grouping_level_new_issues.py

@@ -106,7 +106,15 @@ def _get_hash_for_parent_level(group: Group, id: int, levels_overview: LevelsOve
             .set_where(_get_group_filters(group))
             .set_limit(1)
         )
-        request = SnubaRequest(dataset="events", app_id="grouping", query=query)
+        request = SnubaRequest(
+            dataset="events",
+            app_id="grouping",
+            query=query,
+            tenant_ids={
+                "referrer": "api.group_hashes_levels.get_hash_for_parent_level",
+                "organization_id": group.project.organization_id,
+            },
+        )
         return_hash: str = get_path(snuba.raw_snql_query(request), "data", 0, "hash")  # type: ignore
         cache.set(cache_key, return_hash)
 
@@ -187,10 +195,14 @@ def _query_snuba(group: Group, id: int, offset=None, limit=None):
     if limit is not None:
         query = query.set_limit(limit)
 
-    request = SnubaRequest(dataset="events", app_id="grouping", query=query)
-    return snuba.raw_snql_query(request, referrer="api.group_hashes_levels.get_level_new_issues")[
-        "data"
-    ]
+    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},
+    )
+    return snuba.raw_snql_query(request, referrer)["data"]
 
 
 def _process_snuba_results(query_res, group: Group, id: int, user):

+ 9 - 3
src/sentry/api/endpoints/grouping_levels.py

@@ -112,7 +112,7 @@ class LevelsOverview:
     num_levels: int
 
 
-def get_levels_overview(group):
+def get_levels_overview(group: Group):
     query = (
         Query(Entity("events"))
         .set_select(
@@ -127,8 +127,14 @@ def get_levels_overview(group):
         .set_where(_get_group_filters(group))
         .set_groupby([Column("primary_hash")])
     )
-    request = SnubaRequest(dataset="events", app_id="grouping", query=query)
-    res = snuba.raw_snql_query(request, referrer="api.group_hashes_levels.get_levels_overview")
+    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},
+    )
+    res = snuba.raw_snql_query(request, referrer)
 
     if not res["data"]:
         raise NoEvents()

+ 3 - 1
src/sentry/api/endpoints/organization_events_span_ops.py

@@ -38,8 +38,10 @@ class OrganizationEventsSpanOpsEndpoint(OrganizationEventsEndpointBase):  # type
                 offset=offset,
                 orderby="-count",
             )
+            referrer = "api.organization-events-span-ops"
             snql_query = builder.get_snql_query()
-            results = raw_snql_query(snql_query, "api.organization-events-span-ops")
+            snql_query.tenant_ids = {"referrer": referrer, "organization_id": organization.id}
+            results = raw_snql_query(snql_query, referrer)
             return [SpanOp(op=row["spans_op"], count=row["count"]) for row in results["data"]]
 
         with self.handle_query_errors():

+ 18 - 4
src/sentry/api/endpoints/project_dynamic_sampling.py

@@ -110,10 +110,16 @@ class ProjectDynamicSamplingDistributionEndpoint(ProjectEndpoint):
                 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),
-            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_PROJECT_STATS.value,
+            SnubaRequest(
+                dataset=Dataset.Discover.value,
+                app_id="default",
+                query=snuba_query,
+                tenant_ids={"referrer": referrer, "organization_id": org_id},
+            ),
+            referrer,
         )
         return builder.process_results(data)["data"]
 
@@ -259,13 +265,21 @@ class ProjectDynamicSamplingDistributionEndpoint(ProjectEndpoint):
                 )
             ]
         )
+
+        referrer = Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_TRANSACTIONS.value
+
         snuba_query = snuba_query.set_groupby(
             snuba_query.groupby + [Column("modulo_num"), Column("contexts.key")]
         )
 
         data = raw_snql_query(
-            SnubaRequest(dataset=Dataset.Discover.value, app_id="default", query=snuba_query),
-            referrer=Referrer.DYNAMIC_SAMPLING_DISTRIBUTION_FETCH_TRANSACTIONS.value,
+            SnubaRequest(
+                dataset=Dataset.Discover.value,
+                app_id="default",
+                query=snuba_query,
+                tenant_ids={"referrer": referrer, "organization_id": project.organization_id},
+            ),
+            referrer,
         )["data"]
         return data
 

+ 1 - 1
src/sentry/api/serializers/models/group.py

@@ -899,7 +899,7 @@ class GroupSerializerSnuba(GroupSerializerBase):
         from sentry.search.snuba.executors import get_search_filter
 
         self.environment_ids = environment_ids
-
+        self.organization_id = organization_id
         # XXX: We copy this logic from `PostgresSnubaQueryExecutor.query`. Ideally we
         # should try and encapsulate this logic, but if you're changing this, change it
         # there as well.

+ 1 - 0
src/sentry/api/serializers/models/group_stream.py

@@ -348,6 +348,7 @@ class StreamGroupSerializerSnuba(GroupSerializerSnuba, GroupStatsMixin):
             snuba_tsdb.get_range,
             environment_ids=environment_ids,
             conditions=conditions,
+            tenant_ids={"organization_id": self.organization_id},
             **query_params,
         )
         if error_issue_ids:

Некоторые файлы не были показаны из-за большого количества измененных файлов