ref

Author: brendanhsentry

src/sentry/billing/platform/services/usage/_outcomes_query.py

@@ -61,7 +61,7 @@ def query_orgs_with_usage(request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageR
 
     offset = request.page_token.offset if request.HasField("page_token") else 0
 
-    snuba_request = _build_orgs_query(
+    snuba_request = _build_distinct_orgs_query(
         start=start,
         end=end,
         categories=categories,
@@ -96,7 +96,9 @@ def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse:
     # (e.g., proto ATTACHMENT=3 vs Relay ATTACHMENT=4). Convert before querying.
     categories = [proto_to_sentry_category(c) for c in request.categories]
 
-    snuba_request = _build_query(org_id, start, end, categories, total_outcomes=_BILLABLE_OUTCOMES)
+    snuba_request = _build_usage_query(
+        org_id, start, end, categories, total_outcomes=_BILLABLE_OUTCOMES
+    )
     result = raw_snql_query(snuba_request, referrer=_REFERRER)
     rows = result["data"]
 
@@ -114,7 +116,7 @@ def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse:
     return _build_response(rows)
 
 
-def _build_orgs_query(
+def _build_distinct_orgs_query(
     start: datetime,
     end: datetime,
     categories: Sequence[int],
@@ -150,8 +152,8 @@ def _build_orgs_query(
     )
 
 
-def _build_query(
-    org_id: int,
+def _build_usage_query(
+    org_id: int | None,
     start: datetime,
     end: datetime,
     categories: Sequence[int],
@@ -160,14 +162,16 @@ def _build_query(
     limit: int = _QUERY_LIMIT,
     offset: int = 0,
 ) -> Request:
+    """Build a per-category, per-day usage breakdown query."""
     # Half-open interval [start, end) — standard sentry.snuba.outcomes convention.
     # `end` has already been shifted +1 day in query_outcomes_usage() to convert
     # the proto's inclusive end into the exclusive boundary Snuba expects.
     where = [
         Condition(Column("timestamp"), Op.GTE, start),
         Condition(Column("timestamp"), Op.LT, end),
-        Condition(Column("org_id"), Op.EQ, org_id),
     ]
+    if org_id is not None:
+        where.append(Condition(Column("org_id"), Op.EQ, org_id))
     if categories:
         where.append(Condition(Column("category"), Op.IN, categories))
 
@@ -224,22 +228,31 @@ def _build_query(
         ),
     ]
 
+    groupby = [Column("category"), Column("time")]
+    if org_id is None:
+        select.insert(0, Column("org_id"))
+        groupby.insert(0, Column("org_id"))
+
     query = Query(
         match=Entity("outcomes"),
         select=select,
-        groupby=[Column("category"), Column("time")],
+        groupby=groupby,
         where=where,
         orderby=[OrderBy(Column("time"), Direction.ASC)],
         granularity=Granularity(_DAILY_GRANULARITY),
         limit=Limit(limit),
         offset=Offset(offset),
     )
 
+    tenant_ids: dict[str, int] = {}
+    if org_id is not None:
+        tenant_ids["organization_id"] = org_id
+
     return Request(
         dataset=_DATASET,
         app_id=_APP_ID,
         query=query,
-        tenant_ids={"organization_id": org_id},
+        tenant_ids=tenant_ids,
     )
 
 

tests/sentry/billing/platform/services/usage/test_outcomes_query.py

@@ -18,8 +18,8 @@
 
 from sentry.billing.platform.services.usage._outcomes_query import (
     _BILLABLE_OUTCOMES,
-    _build_query,
     _build_response,
+    _build_usage_query,
     _over_quota_condition,
     _total_function,
     query_orgs_with_usage,
@@ -182,7 +182,7 @@ def test_build_query_with_categories(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(org_id=1, start=start, end=end, categories=[1, 2])
+        snuba_request = _build_usage_query(org_id=1, start=start, end=end, categories=[1, 2])
 
         query = snuba_request.query
         category_conditions = [
@@ -196,7 +196,7 @@ def test_build_query_no_categories(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(org_id=1, start=start, end=end, categories=[])
+        snuba_request = _build_usage_query(org_id=1, start=start, end=end, categories=[])
 
         query = snuba_request.query
         category_conditions = [
@@ -208,7 +208,7 @@ def test_build_query_basic_structure(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(org_id=42, start=start, end=end, categories=[])
+        snuba_request = _build_usage_query(org_id=42, start=start, end=end, categories=[])
 
         assert snuba_request.dataset == "outcomes"
         assert snuba_request.app_id == "billing"
@@ -223,7 +223,7 @@ def test_build_query_groups_by_category_and_time_only(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(org_id=1, start=start, end=end, categories=[])
+        snuba_request = _build_usage_query(org_id=1, start=start, end=end, categories=[])
 
         groupby_names = [col.name for col in snuba_request.query.groupby]
         assert groupby_names == ["category", "time"]
@@ -232,7 +232,7 @@ def test_build_query_total_filters_billable_outcomes(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(
+        snuba_request = _build_usage_query(
             org_id=1, start=start, end=end, categories=[], total_outcomes=_BILLABLE_OUTCOMES
         )
 
@@ -258,7 +258,7 @@ def test_build_query_total_all_outcomes_when_none(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(
+        snuba_request = _build_usage_query(
             org_id=1, start=start, end=end, categories=[], total_outcomes=None
         )
 
@@ -281,7 +281,7 @@ def test_build_query_select_has_sumif_columns(self):
         start = datetime(2025, 3, 1, tzinfo=timezone.utc)
         end = datetime(2025, 3, 31, tzinfo=timezone.utc)
 
-        snuba_request = _build_query(org_id=1, start=start, end=end, categories=[])
+        snuba_request = _build_usage_query(org_id=1, start=start, end=end, categories=[])
 
         select = snuba_request.query.select
         aliases = []