From 8def4d44b1cb62e5f1c2fa05c78bb9bde2384964 Mon Sep 17 00:00:00 2001 From: Brendan Hy Date: Fri, 3 Apr 2026 11:36:56 -0700 Subject: [PATCH 1/5] feat(billing): Add get orgs with usage endpoint --- .../services/usage/_outcomes_query.py | 169 ++++++++++++------ .../platform/services/usage/service.py | 9 +- .../services/usage/test_outcomes_query.py | 143 +++++++++++++++ 3 files changed, 263 insertions(+), 58 deletions(-) diff --git a/src/sentry/billing/platform/services/usage/_outcomes_query.py b/src/sentry/billing/platform/services/usage/_outcomes_query.py index 5e16ef64bab1a2..6a517805850dab 100644 --- a/src/sentry/billing/platform/services/usage/_outcomes_query.py +++ b/src/sentry/billing/platform/services/usage/_outcomes_query.py @@ -7,6 +7,11 @@ from google.protobuf.timestamp_pb2 import Timestamp from sentry_protos.billing.v1.date_pb2 import Date +from sentry_protos.billing.v1.services.usage.v1.endpoint_orgs_with_usage_pb2 import ( + GetOrgsWithUsageRequest, + GetOrgsWithUsageResponse, + PageToken, +) from sentry_protos.billing.v1.services.usage.v1.endpoint_usage_pb2 import ( CategoryUsage, DailyUsage, @@ -21,6 +26,7 @@ Function, Granularity, Limit, + Offset, Op, OrderBy, Query, @@ -48,6 +54,39 @@ _BILLABLE_OUTCOMES = [Outcome.ACCEPTED, Outcome.FILTERED, Outcome.RATE_LIMITED] +def query_orgs_with_usage(request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageResponse: + start = _timestamp_to_datetime(request.start) + end = _timestamp_to_datetime(request.end) + timedelta(days=1) + categories = [proto_to_relay_category(c) for c in request.categories] + + limit = request.limit if request.HasField("limit") else _QUERY_LIMIT + offset = request.page_token.offset if request.HasField("page_token") else 0 + + snuba_request = _build_query( + org_id=None, + start=start, + end=end, + categories=categories, + total_outcomes=_BILLABLE_OUTCOMES, + limit=limit + 1, + offset=offset, + ) + result = raw_snql_query(snuba_request, referrer=_REFERRER) + rows = result["data"] + + has_more = len(rows) > limit + if has_more: + rows.pop() + + response = GetOrgsWithUsageResponse( + organization_ids=[int(row["org_id"]) for row in rows], + ) + if has_more: + response.page_token.CopyFrom(PageToken(offset=offset + limit)) + + return response + + def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse: org_id = request.organization_id start = _timestamp_to_datetime(request.start) @@ -79,89 +118,105 @@ def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse: def _build_query( - org_id: int, + org_id: int | None, start: datetime, end: datetime, categories: Sequence[int], *, total_outcomes: Sequence[int] | None = None, + limit: int = _QUERY_LIMIT, + offset: int = 0, ) -> Request: # 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("org_id"), Op.EQ, org_id), Condition(Column("timestamp"), Op.GTE, start), Condition(Column("timestamp"), Op.LT, end), ] + 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)) + select = [ + Column("category"), + Column("time"), + _total_function(total_outcomes), + Function( + "sumIf", + [Column("quantity"), Function("equals", [Column("outcome"), Outcome.ACCEPTED])], + "accepted", + ), + Function( + "sumIf", + [ + Column("quantity"), + Function("equals", [Column("outcome"), Outcome.RATE_LIMITED]), + ], + "dropped", + ), + Function( + "sumIf", + [Column("quantity"), Function("equals", [Column("outcome"), Outcome.FILTERED])], + "filtered", + ), + Function("sumIf", [Column("quantity"), _over_quota_condition()], "over_quota"), + Function( + "sumIf", + [ + Column("quantity"), + Function( + "and", + [ + Function("equals", [Column("outcome"), Outcome.RATE_LIMITED]), + Function("equals", [Column("reason"), "smart_rate_limit"]), + ], + ), + ], + "spike_protection", + ), + Function( + "sumIf", + [ + Column("quantity"), + Function( + "and", + [ + Function("equals", [Column("outcome"), Outcome.FILTERED]), + Function("startsWith", [Column("reason"), "Sampled:"]), + ], + ), + ], + "dynamic_sampling", + ), + ] + + 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=[ - Column("category"), - Column("time"), - _total_function(total_outcomes), - Function( - "sumIf", - [Column("quantity"), Function("equals", [Column("outcome"), Outcome.ACCEPTED])], - "accepted", - ), - Function( - "sumIf", - [ - Column("quantity"), - Function("equals", [Column("outcome"), Outcome.RATE_LIMITED]), - ], - "dropped", - ), - Function( - "sumIf", - [Column("quantity"), Function("equals", [Column("outcome"), Outcome.FILTERED])], - "filtered", - ), - Function("sumIf", [Column("quantity"), _over_quota_condition()], "over_quota"), - Function( - "sumIf", - [ - Column("quantity"), - Function( - "and", - [ - Function("equals", [Column("outcome"), Outcome.RATE_LIMITED]), - Function("equals", [Column("reason"), "smart_rate_limit"]), - ], - ), - ], - "spike_protection", - ), - Function( - "sumIf", - [ - Column("quantity"), - Function( - "and", - [ - Function("equals", [Column("outcome"), Outcome.FILTERED]), - Function("startsWith", [Column("reason"), "Sampled:"]), - ], - ), - ], - "dynamic_sampling", - ), - ], - groupby=[Column("category"), Column("time")], + select=select, + groupby=groupby, where=where, orderby=[OrderBy(Column("time"), Direction.ASC)], granularity=Granularity(_DAILY_GRANULARITY), - limit=Limit(_QUERY_LIMIT), + 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, ) diff --git a/src/sentry/billing/platform/services/usage/service.py b/src/sentry/billing/platform/services/usage/service.py index cf2452deda1d4e..705c37f589ae0e 100644 --- a/src/sentry/billing/platform/services/usage/service.py +++ b/src/sentry/billing/platform/services/usage/service.py @@ -6,7 +6,10 @@ ) from sentry.billing.platform.core import BillingService, service_method -from sentry.billing.platform.services.usage._outcomes_query import query_outcomes_usage +from sentry.billing.platform.services.usage._outcomes_query import ( + query_orgs_with_usage, + query_outcomes_usage, +) class UsageService(BillingService): @@ -20,3 +23,7 @@ def get_usage(self, request: GetUsageRequest) -> GetUsageResponse: dynamic_sampling. """ return query_outcomes_usage(request) + + @service_method + def get_orgs_with_usage(self, request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageResponse: + return query_orgs_with_usage(request) diff --git a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py b/tests/sentry/billing/platform/services/usage/test_outcomes_query.py index 697a2128ea57b4..994f1b91b13903 100644 --- a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py +++ b/tests/sentry/billing/platform/services/usage/test_outcomes_query.py @@ -5,6 +5,11 @@ from google.protobuf.timestamp_pb2 import Timestamp from sentry_protos.billing.v1.date_pb2 import Date +from sentry_protos.billing.v1.services.usage.v1.endpoint_orgs_with_usage_pb2 import ( + GetOrgsWithUsageRequest, + GetOrgsWithUsageResponse, + PageToken, +) from sentry_protos.billing.v1.services.usage.v1.endpoint_usage_pb2 import ( GetUsageRequest, GetUsageResponse, @@ -13,10 +18,12 @@ from sentry.billing.platform.services.usage._outcomes_query import ( _BILLABLE_OUTCOMES, + _QUERY_LIMIT, _build_query, _build_response, _over_quota_condition, _total_function, + query_orgs_with_usage, query_outcomes_usage, ) from sentry.utils.outcomes import Outcome @@ -363,3 +370,139 @@ def test_query_returns_response(self, mock_query): assert len(response.days) == 1 assert response.days[0].date == Date(year=2025, month=3, day=15) assert response.days[0].usage[0].data.accepted == 200 + + +class TestQueryOrgsWithUsage: + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_returns_org_ids(self, mock_query): + mock_query.return_value = { + "data": [ + {"org_id": 1, **_make_row()}, + {"org_id": 42, **_make_row()}, + {"org_id": 99, **_make_row()}, + ] + } + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end) + + response = query_orgs_with_usage(request) + + assert isinstance(response, GetOrgsWithUsageResponse) + assert list(response.organization_ids) == [1, 42, 99] + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_empty_results(self, mock_query): + mock_query.return_value = {"data": []} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end) + + response = query_orgs_with_usage(request) + + assert list(response.organization_ids) == [] + assert not response.HasField("page_token") + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_query_has_no_org_id_filter(self, mock_query): + mock_query.return_value = {"data": []} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end) + + query_orgs_with_usage(request) + + snuba_request = mock_query.call_args[0][0] + org_conditions = [ + c for c in snuba_request.query.where if hasattr(c, "lhs") and c.lhs.name == "org_id" + ] + assert len(org_conditions) == 0 + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_pagination_no_more_results(self, mock_query): + mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(5)]} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=10) + + response = query_orgs_with_usage(request) + + assert len(response.organization_ids) == 5 + assert not response.HasField("page_token") + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_pagination_has_more(self, mock_query): + mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(11)]} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=10) + + response = query_orgs_with_usage(request) + + assert len(response.organization_ids) == 10 + assert response.HasField("page_token") + assert response.page_token.offset == 10 + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_pagination_with_page_token(self, mock_query): + """page_token offset is forwarded to the Snuba query.""" + mock_query.return_value = {"data": [{"org_id": 99, **_make_row()}]} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest( + start=start, end=end, limit=10, page_token=PageToken(offset=50) + ) + + response = query_orgs_with_usage(request) + + snuba_request = mock_query.call_args[0][0] + assert snuba_request.query.offset.offset == 50 + assert list(response.organization_ids) == [99] + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_pagination_next_offset_accounts_for_current(self, mock_query): + mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(6)]} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest( + start=start, end=end, limit=5, page_token=PageToken(offset=20) + ) + + response = query_orgs_with_usage(request) + + assert len(response.organization_ids) == 5 + assert response.page_token.offset == 25 + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_default_limit_used_when_not_set(self, mock_query): + mock_query.return_value = {"data": []} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end) + + query_orgs_with_usage(request) + + snuba_request = mock_query.call_args[0][0] + assert snuba_request.query.limit.limit == _QUERY_LIMIT + 1 + + @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") + def test_queries_with_limit_plus_one(self, mock_query): + """Snuba query uses limit+1 to detect has_more.""" + mock_query.return_value = {"data": []} + + start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=100) + + query_orgs_with_usage(request) + + snuba_request = mock_query.call_args[0][0] + assert snuba_request.query.limit.limit == 101 From 2c0d4faf9c0cb1cf500f781328cac891846eb069 Mon Sep 17 00:00:00 2001 From: Brendan Hy Date: Fri, 3 Apr 2026 13:52:17 -0700 Subject: [PATCH 2/5] fix import --- src/sentry/billing/platform/services/usage/service.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/sentry/billing/platform/services/usage/service.py b/src/sentry/billing/platform/services/usage/service.py index 705c37f589ae0e..86061587a379b6 100644 --- a/src/sentry/billing/platform/services/usage/service.py +++ b/src/sentry/billing/platform/services/usage/service.py @@ -1,5 +1,9 @@ from __future__ import annotations +from sentry_protos.billing.v1.services.usage.v1.endpoint_orgs_with_usage_pb2 import ( + GetOrgsWithUsageRequest, + GetOrgsWithUsageResponse, +) from sentry_protos.billing.v1.services.usage.v1.endpoint_usage_pb2 import ( GetUsageRequest, GetUsageResponse, From 9f7e56481d38822a6c56f99c8daaf1deac0fed5d Mon Sep 17 00:00:00 2001 From: Brendan Hy Date: Fri, 3 Apr 2026 14:59:07 -0700 Subject: [PATCH 3/5] fix --- .../services/usage/_outcomes_query.py | 65 +++++++++++++------ .../services/usage/test_outcomes_query.py | 42 ++++-------- 2 files changed, 57 insertions(+), 50 deletions(-) diff --git a/src/sentry/billing/platform/services/usage/_outcomes_query.py b/src/sentry/billing/platform/services/usage/_outcomes_query.py index 6a517805850dab..a49a4d826998d9 100644 --- a/src/sentry/billing/platform/services/usage/_outcomes_query.py +++ b/src/sentry/billing/platform/services/usage/_outcomes_query.py @@ -59,22 +59,19 @@ def query_orgs_with_usage(request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageR end = _timestamp_to_datetime(request.end) + timedelta(days=1) categories = [proto_to_relay_category(c) for c in request.categories] - limit = request.limit if request.HasField("limit") else _QUERY_LIMIT offset = request.page_token.offset if request.HasField("page_token") else 0 - snuba_request = _build_query( - org_id=None, + snuba_request = _build_orgs_query( start=start, end=end, categories=categories, - total_outcomes=_BILLABLE_OUTCOMES, - limit=limit + 1, + limit=request.limit + 1, offset=offset, ) result = raw_snql_query(snuba_request, referrer=_REFERRER) rows = result["data"] - has_more = len(rows) > limit + has_more = len(rows) > request.limit if has_more: rows.pop() @@ -82,7 +79,7 @@ def query_orgs_with_usage(request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageR organization_ids=[int(row["org_id"]) for row in rows], ) if has_more: - response.page_token.CopyFrom(PageToken(offset=offset + limit)) + response.page_token.CopyFrom(PageToken(offset=offset + request.limit)) return response @@ -117,8 +114,44 @@ def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse: return _build_response(rows) +def _build_orgs_query( + start: datetime, + end: datetime, + categories: Sequence[int], + *, + limit: int, + offset: int = 0, +) -> Request: + """Build a query that returns distinct org_ids with billable usage.""" + where = [ + Condition(Column("timestamp"), Op.GTE, start), + Condition(Column("timestamp"), Op.LT, end), + Condition(Column("outcome"), Op.IN, list(_BILLABLE_OUTCOMES)), + ] + if categories: + where.append(Condition(Column("category"), Op.IN, categories)) + + query = Query( + match=Entity("outcomes"), + select=[Column("org_id")], + groupby=[Column("org_id")], + where=where, + orderby=[OrderBy(Column("org_id"), Direction.ASC)], + granularity=Granularity(_DAILY_GRANULARITY), + limit=Limit(limit), + offset=Offset(offset), + ) + + return Request( + dataset=_DATASET, + app_id=_APP_ID, + query=query, + tenant_ids={}, + ) + + def _build_query( - org_id: int | None, + org_id: int, start: datetime, end: datetime, categories: Sequence[int], @@ -133,9 +166,8 @@ def _build_query( 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)) @@ -192,15 +224,10 @@ 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=groupby, + groupby=[Column("category"), Column("time")], where=where, orderby=[OrderBy(Column("time"), Direction.ASC)], granularity=Granularity(_DAILY_GRANULARITY), @@ -208,15 +235,11 @@ def _build_query( 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=tenant_ids, + tenant_ids={"organization_id": org_id}, ) diff --git a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py b/tests/sentry/billing/platform/services/usage/test_outcomes_query.py index 994f1b91b13903..3a2c8731eb0e8c 100644 --- a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py +++ b/tests/sentry/billing/platform/services/usage/test_outcomes_query.py @@ -18,7 +18,6 @@ from sentry.billing.platform.services.usage._outcomes_query import ( _BILLABLE_OUTCOMES, - _QUERY_LIMIT, _build_query, _build_response, _over_quota_condition, @@ -377,15 +376,15 @@ class TestQueryOrgsWithUsage: def test_returns_org_ids(self, mock_query): mock_query.return_value = { "data": [ - {"org_id": 1, **_make_row()}, - {"org_id": 42, **_make_row()}, - {"org_id": 99, **_make_row()}, + {"org_id": 1}, + {"org_id": 42}, + {"org_id": 99}, ] } start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) - request = GetOrgsWithUsageRequest(start=start, end=end) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=100) response = query_orgs_with_usage(request) @@ -398,7 +397,7 @@ def test_empty_results(self, mock_query): start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) - request = GetOrgsWithUsageRequest(start=start, end=end) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=100) response = query_orgs_with_usage(request) @@ -406,24 +405,22 @@ def test_empty_results(self, mock_query): assert not response.HasField("page_token") @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") - def test_query_has_no_org_id_filter(self, mock_query): + def test_query_groups_only_by_org_id(self, mock_query): mock_query.return_value = {"data": []} start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) - request = GetOrgsWithUsageRequest(start=start, end=end) + request = GetOrgsWithUsageRequest(start=start, end=end, limit=100) query_orgs_with_usage(request) snuba_request = mock_query.call_args[0][0] - org_conditions = [ - c for c in snuba_request.query.where if hasattr(c, "lhs") and c.lhs.name == "org_id" - ] - assert len(org_conditions) == 0 + groupby_names = [col.name for col in snuba_request.query.groupby] + assert groupby_names == ["org_id"] @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") def test_pagination_no_more_results(self, mock_query): - mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(5)]} + mock_query.return_value = {"data": [{"org_id": i} for i in range(5)]} start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) @@ -436,7 +433,7 @@ def test_pagination_no_more_results(self, mock_query): @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") def test_pagination_has_more(self, mock_query): - mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(11)]} + mock_query.return_value = {"data": [{"org_id": i} for i in range(11)]} start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) @@ -451,7 +448,7 @@ def test_pagination_has_more(self, mock_query): @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") def test_pagination_with_page_token(self, mock_query): """page_token offset is forwarded to the Snuba query.""" - mock_query.return_value = {"data": [{"org_id": 99, **_make_row()}]} + mock_query.return_value = {"data": [{"org_id": 99}]} start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) @@ -467,7 +464,7 @@ def test_pagination_with_page_token(self, mock_query): @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") def test_pagination_next_offset_accounts_for_current(self, mock_query): - mock_query.return_value = {"data": [{"org_id": i, **_make_row()} for i in range(6)]} + mock_query.return_value = {"data": [{"org_id": i} for i in range(6)]} start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) @@ -480,19 +477,6 @@ def test_pagination_next_offset_accounts_for_current(self, mock_query): assert len(response.organization_ids) == 5 assert response.page_token.offset == 25 - @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") - def test_default_limit_used_when_not_set(self, mock_query): - mock_query.return_value = {"data": []} - - start = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) - end = _make_timestamp(datetime(2025, 3, 1, tzinfo=timezone.utc)) - request = GetOrgsWithUsageRequest(start=start, end=end) - - query_orgs_with_usage(request) - - snuba_request = mock_query.call_args[0][0] - assert snuba_request.query.limit.limit == _QUERY_LIMIT + 1 - @patch("sentry.billing.platform.services.usage._outcomes_query.raw_snql_query") def test_queries_with_limit_plus_one(self, mock_query): """Snuba query uses limit+1 to detect has_more.""" From 92dfd12a430e50490e27e4742d4f25e8172a94b3 Mon Sep 17 00:00:00 2001 From: Brendan Hy Date: Fri, 3 Apr 2026 15:02:09 -0700 Subject: [PATCH 4/5] merge conflict --- src/sentry/billing/platform/services/usage/_outcomes_query.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sentry/billing/platform/services/usage/_outcomes_query.py b/src/sentry/billing/platform/services/usage/_outcomes_query.py index 042bf6f0504a9e..be2d5ac64bee78 100644 --- a/src/sentry/billing/platform/services/usage/_outcomes_query.py +++ b/src/sentry/billing/platform/services/usage/_outcomes_query.py @@ -57,7 +57,7 @@ def query_orgs_with_usage(request: GetOrgsWithUsageRequest) -> GetOrgsWithUsageResponse: start = _timestamp_to_datetime(request.start) end = _timestamp_to_datetime(request.end) + timedelta(days=1) - categories = [proto_to_relay_category(c) for c in request.categories] + categories = [proto_to_sentry_category(c) for c in request.categories] offset = request.page_token.offset if request.HasField("page_token") else 0 From 927237cbf8700edb7ca7bf9688e12fb82a228597 Mon Sep 17 00:00:00 2001 From: Brendan Hy Date: Fri, 3 Apr 2026 15:31:26 -0700 Subject: [PATCH 5/5] ref --- .../services/usage/_outcomes_query.py | 29 ++++++++++++++----- .../services/usage/test_outcomes_query.py | 16 +++++----- 2 files changed, 29 insertions(+), 16 deletions(-) diff --git a/src/sentry/billing/platform/services/usage/_outcomes_query.py b/src/sentry/billing/platform/services/usage/_outcomes_query.py index be2d5ac64bee78..68c6080652e81b 100644 --- a/src/sentry/billing/platform/services/usage/_outcomes_query.py +++ b/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,10 +228,15 @@ 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), @@ -235,11 +244,15 @@ def _build_query( 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, ) diff --git a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py b/tests/sentry/billing/platform/services/usage/test_outcomes_query.py index 3a2c8731eb0e8c..e1af501acc60c5 100644 --- a/tests/sentry/billing/platform/services/usage/test_outcomes_query.py +++ b/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 = []