diff --git a/src/sentry/billing/__init__.py b/src/sentry/billing/__init__.py new file mode 100644 index 00000000000000..2ea6afb085a317 --- /dev/null +++ b/src/sentry/billing/__init__.py @@ -0,0 +1,9 @@ +from sentry.billing.interface import BillingService +from sentry.billing.sentry import SentryUsageTrackingService + + +class Service(BillingService): + usage_tracking = SentryUsageTrackingService() + + +billing_service = Service() diff --git a/src/sentry/billing/config.py b/src/sentry/billing/config.py new file mode 100644 index 00000000000000..22eec466d6f8a3 --- /dev/null +++ b/src/sentry/billing/config.py @@ -0,0 +1,33 @@ +from enum import IntEnum + +from sentry.constants import DataCategory +from sentry.utils.outcomes import Outcome + + +def _id_for(data_category: DataCategory, outcome: Outcome) -> int: + return data_category | (outcome << 10) + + +# Usage category represents actual tracked usage within an organization. +# Importantly, billed usage is separately derived from tracked usage and just +# because usage is tracked does not mean it is billed. +# +# For example: ERROR_FILTERED is tracked but not billed. +class UsageCategoryId(IntEnum): + ERROR_ACCEPTED = _id_for(DataCategory.ERROR, Outcome.ACCEPTED) + ERROR_FILTERED = _id_for(DataCategory.ERROR, Outcome.FILTERED) + ERROR_RATE_LIMITED = _id_for(DataCategory.ERROR, Outcome.RATE_LIMITED) + ERROR_INVALID = _id_for(DataCategory.ERROR, Outcome.INVALID) + ERROR_ABUSE = _id_for(DataCategory.ERROR, Outcome.ABUSE) + ERROR_CLIENT_DISCARD = _id_for(DataCategory.ERROR, Outcome.CLIENT_DISCARD) + ERROR_CARDINALITY_LIMITED = _id_for(DataCategory.ERROR, Outcome.CARDINALITY_LIMITED) + # TODO: TRANSACTION... + + def data_category(self) -> DataCategory: + return DataCategory(self.value & 0x3FF) + + def outcome(self) -> Outcome: + return Outcome(self.value >> 10) + + def api_name(self) -> str: + return f"{self.data_category().api_name()}_{self.outcome().api_name()}" diff --git a/src/sentry/billing/interface.py b/src/sentry/billing/interface.py new file mode 100644 index 00000000000000..2d40931226777a --- /dev/null +++ b/src/sentry/billing/interface.py @@ -0,0 +1,13 @@ +from typing import Protocol + +from sentry.billing.usage import UsageTrackingService + + +class BillingService(Protocol): + usage_tracking: UsageTrackingService + # plan_management: PlanManagementService + # customer_billing: CustomerBillingService + # quota_management: QuotaManagementService + # committed_spend: CommittedSpendService + # billing_calculation: BillingCalculationService + # invoicing: InvoicingService diff --git a/src/sentry/billing/sentry.py b/src/sentry/billing/sentry.py new file mode 100644 index 00000000000000..1316e85c55aec1 --- /dev/null +++ b/src/sentry/billing/sentry.py @@ -0,0 +1,63 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Any, Literal + +from sentry.billing.config import UsageCategoryId +from sentry.billing.usage import UsageProperties, UsageTrackingService +from sentry.snuba.outcomes import QueryDefinition, run_outcomes_query_timeseries +from sentry.utils.outcomes import track_outcome + + +class SentryUsageTrackingService(UsageTrackingService): + def record_usage( + self, + org_id: int, + usage_category_id: UsageCategoryId, + properties: UsageProperties, + timestamp: datetime | None = None, + ) -> None: + track_outcome( + org_id=org_id, + project_id=properties["project_id"], + key_id=properties.get("key_id"), + outcome=usage_category_id.outcome(), + reason=properties.get("reason"), + timestamp=timestamp, + event_id=properties.get("event_id"), + category=usage_category_id.data_category(), + quantity=properties.get("quantity", 1), + ) + + def get_aggregated_usage( + self, + org_id: int, + usage_category_ids: list[UsageCategoryId], + start: datetime, + end: datetime, + filter_properties: UsageProperties | None = None, + group_by: list[str] | None = None, + values: list[str] | None = None, + window_size: Literal["1h", "1d"] = "1d", + ) -> Mapping[UsageCategoryId, list[dict[str, Any]]]: + # TODO: Can we query multiple category/outcome at a time + return { + usage_category_id: run_outcomes_query_timeseries( + QueryDefinition( + fields=values or [], + start=start.isoformat(), + end=end.isoformat(), + organization_id=org_id, + project_ids=( + [filter_properties["project_id"]] if filter_properties is not None else None + ), + key_id=filter_properties["key_id"] if filter_properties is not None else None, + interval=window_size, + outcome=[usage_category_id.outcome().api_name()], + group_by=group_by or [], + category=[usage_category_id.data_category().api_name()], + reason=filter_properties["reason"] if filter_properties is not None else None, + ), + tenant_ids={"organization_id": org_id}, + ) + for usage_category_id in usage_category_ids + } diff --git a/src/sentry/billing/usage.py b/src/sentry/billing/usage.py new file mode 100644 index 00000000000000..68ec8ba7b1f0a8 --- /dev/null +++ b/src/sentry/billing/usage.py @@ -0,0 +1,40 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Any, Literal, NotRequired, Protocol, TypedDict + +from sentry.billing.config import UsageCategoryId + + +class UsageProperties(TypedDict): + project_id: int + event_id: NotRequired[str] + key_id: NotRequired[int] + reason: NotRequired[str] + quantity: NotRequired[int] + idempotency_key: NotRequired[str] # TODO: seems useful? + + +class UsageTrackingService(Protocol): + """Service for recording and querying usage data""" + + def record_usage( + self, + org_id: int, + usage_category_id: UsageCategoryId, + properties: UsageProperties, + timestamp: datetime | None = None, + ) -> None: + """Record usage for a specific usage category""" + + def get_aggregated_usage( + self, + org_id: int, + usage_category_ids: list[UsageCategoryId], + start: datetime, + end: datetime, + filter_properties: UsageProperties | None = None, # TODO: filter-specific type + group_by: list[str] | None = None, + values: list[str] | None = None, + window_size: Literal["1h", "1d"] = "1d", + ) -> Mapping[UsageCategoryId, list[dict[str, Any]]]: # TODO: typing as AggregatedUsage + """Retrieve aggregated usage data across multiple usage categories""" diff --git a/tests/sentry/billing/test_sentry.py b/tests/sentry/billing/test_sentry.py new file mode 100644 index 00000000000000..d0ba4cc48c03ee --- /dev/null +++ b/tests/sentry/billing/test_sentry.py @@ -0,0 +1,101 @@ +from datetime import timedelta + +from django.utils import timezone + +from sentry.billing.config import UsageCategoryId +from sentry.billing.sentry import SentryUsageTrackingService +from sentry.testutils.cases import TestCase + + +class TestSentryUsageTrackingService(TestCase): + def setUp(self): + self.service = SentryUsageTrackingService() + + self.organization_id = 1 + self.project_id = 2 + + self.timestamp = timezone.now() - timedelta(hours=1) + self.before_timestamp = self.timestamp - timedelta(minutes=30) + self.after_timestamp = self.timestamp + timedelta(minutes=30) + + def test_record_and_read_basic_usage(self): + self.service.record_usage( + org_id=self.organization_id, + usage_category_id=UsageCategoryId.ERROR_ACCEPTED, + properties={"project_id": self.project_id, "quantity": 3}, + timestamp=self.timestamp, + ) + + result = self.service.get_aggregated_usage( + org_id=self.organization_id, + usage_category_ids=[UsageCategoryId.ERROR_ACCEPTED], + start=self.before_timestamp, + end=self.after_timestamp, + values=["sum(quantity)"], + ) + + # TODO: Verify actual results, not shape + assert isinstance(result, dict) + assert UsageCategoryId.ERROR_ACCEPTED in result + assert isinstance(result[UsageCategoryId.ERROR_ACCEPTED], list) + + def test_record_and_read_with_filtering(self): + key_id = 12345 + timestamp = timezone.now() - timedelta(hours=1) + + self.service.record_usage( + org_id=self.organization_id, + usage_category_id=UsageCategoryId.ERROR_FILTERED, + properties={ + "project_id": self.project_id, + "key_id": key_id, + "reason": "test_filter", + "quantity": 2, + }, + timestamp=timestamp, + ) + + result = self.service.get_aggregated_usage( + org_id=self.organization_id, + usage_category_ids=[UsageCategoryId.ERROR_FILTERED], + start=self.before_timestamp, + end=self.after_timestamp, + filter_properties={ + "project_id": self.project_id, + "key_id": key_id, + "reason": "test_filter", + }, + values=["sum(quantity)"], + ) + + # TODO: Verify actual results, not shape + assert isinstance(result, dict) + assert UsageCategoryId.ERROR_FILTERED in result + + def test_record_and_read_multiple_categories(self): + timestamp = timezone.now() - timedelta(hours=1) + + categories_to_record = [UsageCategoryId.ERROR_ACCEPTED, UsageCategoryId.ERROR_RATE_LIMITED] + + for category in categories_to_record: + self.service.record_usage( + org_id=self.organization_id, + usage_category_id=category, + properties={"project_id": self.project_id}, + timestamp=timestamp, + ) + + result = self.service.get_aggregated_usage( + org_id=self.organization_id, + usage_category_ids=categories_to_record, + start=self.before_timestamp, + end=self.after_timestamp, + values=["sum(quantity)"], + ) + + # TODO: Verify actual results, not shape + assert isinstance(result, dict) + assert len(result) == len(categories_to_record) + for category in categories_to_record: + assert category in result + assert isinstance(result[category], list)