diff --git a/src/sentry/billing/platform/services/quota/__init__.py b/src/sentry/billing/platform/services/quota/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/src/sentry/billing/platform/services/quota/quota_config_mapping.py b/src/sentry/billing/platform/services/quota/quota_config_mapping.py new file mode 100644 index 000000000000..38de29084e6f --- /dev/null +++ b/src/sentry/billing/platform/services/quota/quota_config_mapping.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import logging + +from sentry_protos.billing.v1.quota_config_pb2 import QuotaConfig as ProtoQuotaConfig +from sentry_protos.billing.v1.quota_config_pb2 import QuotaScope as ProtoQuotaScope + +from sentry.billing.platform.services.category_mapping import ( + proto_to_sentry_category, + sentry_to_proto_category, +) +from sentry.constants import DataCategory +from sentry.quotas.base import QuotaConfig, QuotaScope + +logger = logging.getLogger(__name__) + + +def sentry_to_proto_quota_config(quota: QuotaConfig) -> ProtoQuotaConfig: + """Convert a sentry QuotaConfig to its proto equivalent.""" + kwargs: dict = { + "id": quota.id or "", + "categories": sorted(sentry_to_proto_category(int(c)) for c in quota.categories), + "scope": ProtoQuotaScope.ValueType(int(quota.scope)) + if quota.scope is not None + else ProtoQuotaScope.QUOTA_SCOPE_UNSPECIFIED, + } + if quota.scope_id is not None: + kwargs["scope_id"] = quota.scope_id + if quota.limit is not None: + kwargs["limit"] = quota.limit + if quota.window is not None: + kwargs["window"] = quota.window + if quota.reason_code is not None: + kwargs["reason_code"] = quota.reason_code + + return ProtoQuotaConfig(**kwargs) + + +def proto_to_sentry_quota_config(proto_quota: ProtoQuotaConfig) -> QuotaConfig | None: + """Convert a proto QuotaConfig to its sentry equivalent. + + Returns None if the proto had categories but none could be mapped, + since an empty category set would broaden the quota to all data. + """ + categories: list[DataCategory] = [] + for c in proto_quota.categories: + sentry_cat = proto_to_sentry_category(c) + try: + categories.append(DataCategory(sentry_cat)) + except ValueError: + logger.error( + "quota_config_mapping.unknown_category", + extra={"proto_category": c, "mapped_value": sentry_cat}, + ) + + if proto_quota.categories and not categories: + logger.warning( + "quota_config_mapping.all_categories_dropped", + extra={"proto_categories": list(proto_quota.categories)}, + ) + return None + + try: + scope = ( + QuotaScope(proto_quota.scope) + if proto_quota.scope != ProtoQuotaScope.QUOTA_SCOPE_UNSPECIFIED + else QuotaScope.ORGANIZATION + ) + except ValueError: + logger.error( + "quota_config_mapping.unknown_scope", + extra={"proto_scope": proto_quota.scope}, + ) + scope = QuotaScope.ORGANIZATION + + return QuotaConfig( + id=proto_quota.id or None, + categories=categories, + scope=scope, + scope_id=proto_quota.scope_id if proto_quota.HasField("scope_id") else None, + limit=proto_quota.limit if proto_quota.HasField("limit") else None, + window=proto_quota.window if proto_quota.HasField("window") else None, + reason_code=proto_quota.reason_code if proto_quota.HasField("reason_code") else None, + ) diff --git a/tests/sentry/billing/platform/services/quota/__init__.py b/tests/sentry/billing/platform/services/quota/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/sentry/billing/platform/services/quota/test_quota_config_mapping.py b/tests/sentry/billing/platform/services/quota/test_quota_config_mapping.py new file mode 100644 index 000000000000..e058dd66cfd3 --- /dev/null +++ b/tests/sentry/billing/platform/services/quota/test_quota_config_mapping.py @@ -0,0 +1,228 @@ +from sentry_protos.billing.v1.data_category_pb2 import DataCategory as ProtoDataCategory +from sentry_protos.billing.v1.quota_config_pb2 import QuotaConfig as ProtoQuotaConfig +from sentry_protos.billing.v1.quota_config_pb2 import QuotaScope as ProtoQuotaScope + +from sentry.billing.platform.services.quota.quota_config_mapping import ( + proto_to_sentry_quota_config, + sentry_to_proto_quota_config, +) +from sentry.constants import DataCategory +from sentry.quotas.base import QuotaConfig, QuotaScope + + +class TestSentryToProtoQuotaConfig: + def test_basic_conversion(self): + quota = QuotaConfig( + id="test_quota", + categories=[DataCategory.ERROR, DataCategory.TRANSACTION], + scope=QuotaScope.ORGANIZATION, + limit=1000, + window=3600, + reason_code="org_limit", + ) + proto = sentry_to_proto_quota_config(quota) + + assert proto.id == "test_quota" + assert proto.scope == ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION + assert proto.limit == 1000 + assert proto.window == 3600 + assert proto.reason_code == "org_limit" + assert set(proto.categories) == { + ProtoDataCategory.DATA_CATEGORY_ERROR, + ProtoDataCategory.DATA_CATEGORY_TRANSACTION, + } + + def test_reject_all_quota(self): + quota = QuotaConfig( + id=None, + categories=[DataCategory.ERROR], + scope=QuotaScope.ORGANIZATION, + limit=0, + reason_code="blocked", + ) + proto = sentry_to_proto_quota_config(quota) + + assert proto.id == "" + assert proto.limit == 0 + + def test_scope_id_included(self): + quota = QuotaConfig( + id="q", + categories=[DataCategory.ERROR], + scope=QuotaScope.PROJECT, + scope_id="123", + limit=100, + window=60, + reason_code="proj_limit", + ) + proto = sentry_to_proto_quota_config(quota) + assert proto.scope_id == "123" + + +class TestProtoToSentryQuotaConfig: + def test_basic_conversion(self): + proto = ProtoQuotaConfig( + id="test_quota", + categories=[ + ProtoDataCategory.DATA_CATEGORY_ERROR, + ProtoDataCategory.DATA_CATEGORY_TRANSACTION, + ], + scope=ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION, + limit=1000, + window=3600, + reason_code="org_limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.id == "test_quota" + assert result.scope == QuotaScope.ORGANIZATION + assert result.limit == 1000 + assert result.window == 3600 + assert result.reason_code == "org_limit" + assert DataCategory.ERROR in result.categories + assert DataCategory.TRANSACTION in result.categories + + def test_unspecified_scope_defaults_to_organization(self): + proto = ProtoQuotaConfig( + id="q", + categories=[ProtoDataCategory.DATA_CATEGORY_ERROR], + scope=ProtoQuotaScope.QUOTA_SCOPE_UNSPECIFIED, + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.scope == QuotaScope.ORGANIZATION + + def test_unknown_category_skipped(self): + """Unmapped proto category values are skipped instead of raising.""" + unmapped_value = ProtoDataCategory.ValueType(9999) + proto = ProtoQuotaConfig( + id="q", + categories=[ + ProtoDataCategory.DATA_CATEGORY_ERROR, + unmapped_value, + ], + scope=ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION, + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert len(result.categories) == 1 + assert DataCategory.ERROR in result.categories + + def test_all_categories_unmapped_returns_none(self): + """If all categories are unmapped, return None to avoid broadening quota.""" + proto = ProtoQuotaConfig( + id="q", + categories=[ProtoDataCategory.ValueType(9999), ProtoDataCategory.ValueType(9998)], + scope=ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION, + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is None + + def test_empty_categories_preserved(self): + """Proto with no categories should map to empty (all-data quota), not None.""" + proto = ProtoQuotaConfig( + id="q", + scope=ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION, + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert len(result.categories) == 0 + + def test_unknown_scope_falls_back_to_organization(self): + """Unknown scope values fall back to ORGANIZATION instead of raising.""" + proto = ProtoQuotaConfig( + id="q", + categories=[ProtoDataCategory.DATA_CATEGORY_ERROR], + scope=ProtoQuotaScope.ValueType(999), + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.scope == QuotaScope.ORGANIZATION + + def test_project_scope(self): + proto = ProtoQuotaConfig( + id="q", + categories=[ProtoDataCategory.DATA_CATEGORY_ERROR], + scope=ProtoQuotaScope.QUOTA_SCOPE_PROJECT, + scope_id="42", + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.scope == QuotaScope.PROJECT + assert result.scope_id == "42" + + def test_key_scope(self): + proto = ProtoQuotaConfig( + id="q", + categories=[ProtoDataCategory.DATA_CATEGORY_ERROR], + scope=ProtoQuotaScope.QUOTA_SCOPE_KEY, + limit=100, + window=60, + reason_code="limit", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.scope == QuotaScope.KEY + + def test_reject_all_quota(self): + proto = ProtoQuotaConfig( + categories=[ProtoDataCategory.DATA_CATEGORY_ERROR], + scope=ProtoQuotaScope.QUOTA_SCOPE_ORGANIZATION, + limit=0, + reason_code="blocked", + ) + result = proto_to_sentry_quota_config(proto) + + assert result is not None + assert result.id is None + assert result.limit == 0 + + +class TestRoundTrip: + def test_roundtrip_sentry_to_proto_to_sentry(self): + original = QuotaConfig( + id="roundtrip", + categories=[DataCategory.ERROR, DataCategory.SPAN], + scope=QuotaScope.PROJECT, + scope_id="99", + limit=500, + window=60, + reason_code="rate_limit", + ) + proto = sentry_to_proto_quota_config(original) + restored = proto_to_sentry_quota_config(proto) + + assert restored is not None + assert restored.id == original.id + assert restored.categories == original.categories + assert restored.scope == original.scope + assert restored.scope_id == original.scope_id + assert restored.limit == original.limit + assert restored.window == original.window + assert restored.reason_code == original.reason_code