diff --git a/src/sentry/billing/platform/services/usage/_outcomes_query.py b/src/sentry/billing/platform/services/usage/_outcomes_query.py index 8efbec8be28426..68c6080652e81b 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,36 @@ _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_sentry_category(c) for c in request.categories] + + offset = request.page_token.offset if request.HasField("page_token") else 0 + + snuba_request = _build_distinct_orgs_query( + start=start, + end=end, + categories=categories, + limit=request.limit + 1, + offset=offset, + ) + result = raw_snql_query(snuba_request, referrer=_REFERRER) + rows = result["data"] + + has_more = len(rows) > request.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 + request.limit)) + + return response + + def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse: org_id = request.organization_id start = _timestamp_to_datetime(request.start) @@ -60,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"] @@ -78,90 +116,143 @@ def query_outcomes_usage(request: GetUsageRequest) -> GetUsageResponse: return _build_response(rows) -def _build_query( - org_id: int, +def _build_distinct_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_usage_query( + 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: + """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("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..86061587a379b6 100644 --- a/src/sentry/billing/platform/services/usage/service.py +++ b/src/sentry/billing/platform/services/usage/service.py @@ -1,12 +1,19 @@ 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 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 +27,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..e1af501acc60c5 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,11 @@ 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, query_outcomes_usage, ) from sentry.utils.outcomes import Outcome @@ -176,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 = [ @@ -190,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 = [ @@ -202,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" @@ -217,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"] @@ -226,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 ) @@ -252,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 ) @@ -275,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 = [] @@ -363,3 +369,124 @@ 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}, + {"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, limit=100) + + 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, limit=100) + + 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_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, limit=100) + + query_orgs_with_usage(request) + + snuba_request = mock_query.call_args[0][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} 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} 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}]} + + 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} 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_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