1use std::fmt;
2use std::str::FromStr;
3
4use relay_base_schema::metrics::MetricNamespace;
5use relay_base_schema::organization::OrganizationId;
6use relay_base_schema::project::{ProjectId, ProjectKey};
7use serde::{Deserialize, Serialize};
8use smallvec::SmallVec;
9
10#[doc(inline)]
11pub use relay_base_schema::data_category::DataCategory;
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
19pub struct Scoping {
20 pub organization_id: OrganizationId,
22
23 pub project_id: ProjectId,
25
26 pub project_key: ProjectKey,
28
29 pub key_id: Option<u64>,
31}
32
33impl Scoping {
34 pub fn item(&self, category: DataCategory) -> ItemScoping {
40 ItemScoping {
41 category,
42 scoping: *self,
43 namespace: MetricNamespaceScoping::None,
44 }
45 }
46
47 pub fn metric_bucket(&self, namespace: MetricNamespace) -> ItemScoping {
53 ItemScoping {
54 category: DataCategory::MetricBucket,
55 scoping: *self,
56 namespace: MetricNamespaceScoping::Some(namespace),
57 }
58 }
59}
60
61#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash, PartialOrd)]
66pub enum MetricNamespaceScoping {
67 #[default]
71 None,
72
73 Some(MetricNamespace),
75
76 Any,
81}
82
83impl MetricNamespaceScoping {
84 pub fn matches(&self, namespace: MetricNamespace) -> bool {
90 match self {
91 Self::None => false,
92 Self::Some(ns) => *ns == namespace,
93 Self::Any => true,
94 }
95 }
96}
97
98impl From<MetricNamespace> for MetricNamespaceScoping {
99 fn from(namespace: MetricNamespace) -> Self {
100 Self::Some(namespace)
101 }
102}
103
104#[derive(Debug, Copy, Clone, Eq, PartialEq)]
109pub struct ItemScoping {
110 pub category: DataCategory,
112
113 pub scoping: Scoping,
115
116 pub namespace: MetricNamespaceScoping,
118}
119
120impl std::ops::Deref for ItemScoping {
121 type Target = Scoping;
122
123 fn deref(&self) -> &Self::Target {
124 &self.scoping
125 }
126}
127
128impl ItemScoping {
129 pub fn scope_id(&self, scope: QuotaScope) -> Option<u64> {
134 match scope {
135 QuotaScope::Global => None,
136 QuotaScope::Organization => Some(self.organization_id.value()),
137 QuotaScope::Project => Some(self.project_id.value()),
138 QuotaScope::Key => self.key_id,
139 QuotaScope::Unknown => None,
140 }
141 }
142
143 pub(crate) fn matches_categories(&self, categories: &DataCategories) -> bool {
145 categories.is_empty() || categories.contains(&self.category)
150 }
151
152 pub(crate) fn matches_namespaces<'a, I>(&self, namespaces: I) -> bool
166 where
167 I: IntoIterator<Item = &'a MetricNamespace>,
168 {
169 let mut iter = namespaces.into_iter().peekable();
170 iter.peek().is_none() || iter.any(|ns| self.namespace.matches(*ns))
171 }
172}
173
174#[derive(Clone, Copy, Debug, PartialEq, Eq)]
179pub enum CategoryUnit {
180 Count,
182 Bytes,
184 Milliseconds,
186}
187
188impl CategoryUnit {
189 fn from(category: &DataCategory) -> Option<Self> {
190 match category {
191 DataCategory::Default
192 | DataCategory::Error
193 | DataCategory::Transaction
194 | DataCategory::Replay
195 | DataCategory::DoNotUseReplayVideo
196 | DataCategory::Security
197 | DataCategory::Profile
198 | DataCategory::ProfileIndexed
199 | DataCategory::TransactionProcessed
200 | DataCategory::TransactionIndexed
201 | DataCategory::LogItem
202 | DataCategory::Span
203 | DataCategory::SpanIndexed
204 | DataCategory::MonitorSeat
205 | DataCategory::Monitor
206 | DataCategory::MetricBucket
207 | DataCategory::UserReportV2
208 | DataCategory::ProfileChunk
209 | DataCategory::ProfileChunkUi
210 | DataCategory::Uptime
211 | DataCategory::MetricSecond
212 | DataCategory::AttachmentItem
213 | DataCategory::SeerAutofix
214 | DataCategory::SeerScanner
215 | DataCategory::PreventUser
216 | DataCategory::PreventReview
217 | DataCategory::Session
218 | DataCategory::SizeAnalysis
219 | DataCategory::InstallableBuild
220 | DataCategory::TraceMetric => Some(Self::Count),
221 DataCategory::Attachment | DataCategory::LogByte => Some(Self::Bytes),
222 DataCategory::ProfileDuration | DataCategory::ProfileDurationUi => {
223 Some(Self::Milliseconds)
224 }
225
226 DataCategory::Unknown => None,
227 }
228 }
229}
230
231pub type DataCategories = SmallVec<[DataCategory; 8]>;
237
238#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)]
247#[serde(rename_all = "lowercase")]
248pub enum QuotaScope {
249 Global,
251 Organization,
255
256 Project,
260
261 Key,
265
266 #[serde(other)]
268 Unknown,
269}
270
271impl QuotaScope {
272 pub fn from_name(string: &str) -> Self {
276 match string {
277 "global" => Self::Global,
278 "organization" => Self::Organization,
279 "project" => Self::Project,
280 "key" => Self::Key,
281 _ => Self::Unknown,
282 }
283 }
284
285 pub fn name(self) -> &'static str {
289 match self {
290 Self::Global => "global",
291 Self::Key => "key",
292 Self::Project => "project",
293 Self::Organization => "organization",
294 Self::Unknown => "unknown",
295 }
296 }
297}
298
299impl fmt::Display for QuotaScope {
300 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301 write!(f, "{}", self.name())
302 }
303}
304
305impl FromStr for QuotaScope {
306 type Err = ();
307
308 fn from_str(string: &str) -> Result<Self, Self::Err> {
309 Ok(Self::from_name(string))
310 }
311}
312
313fn default_scope() -> QuotaScope {
314 QuotaScope::Organization
315}
316
317#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize, Hash)]
322pub struct ReasonCode(String);
323
324impl ReasonCode {
325 pub fn new<S: Into<String>>(code: S) -> Self {
331 Self(code.into())
332 }
333
334 pub fn as_str(&self) -> &str {
336 &self.0
337 }
338}
339
340impl fmt::Display for ReasonCode {
341 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
342 self.0.fmt(f)
343 }
344}
345
346#[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
359#[serde(rename_all = "camelCase")]
360pub struct Quota {
361 #[serde(default)]
365 pub id: Option<String>,
366
367 #[serde(default = "DataCategories::new")]
371 pub categories: DataCategories,
372
373 #[serde(default = "default_scope")]
378 pub scope: QuotaScope,
379
380 #[serde(default, skip_serializing_if = "Option::is_none")]
385 pub scope_id: Option<String>,
386
387 #[serde(default)]
396 pub limit: Option<u64>,
397
398 #[serde(default, skip_serializing_if = "Option::is_none")]
403 pub window: Option<u64>,
404
405 pub namespace: Option<MetricNamespace>,
409
410 #[serde(default, skip_serializing_if = "Option::is_none")]
415 pub reason_code: Option<ReasonCode>,
416}
417
418impl Quota {
419 pub fn is_valid(&self) -> bool {
426 if self.namespace == Some(MetricNamespace::Unsupported) {
427 return false;
428 }
429
430 let mut units = self.categories.iter().filter_map(CategoryUnit::from);
431
432 match units.next() {
433 None if !self.categories.is_empty() => false,
435 _ if self.limit == Some(0) => true,
437 None => false,
439 Some(unit) => units.all(|u| u == unit),
441 }
442 }
443
444 fn matches_scope(&self, scoping: ItemScoping) -> bool {
451 if self.scope == QuotaScope::Global {
452 return true;
453 }
454
455 let Some(scope_id) = self.scope_id.as_ref() else {
459 return true;
460 };
461
462 let Ok(parsed) = scope_id.parse::<u64>() else {
465 return false;
466 };
467
468 scoping.scope_id(self.scope) == Some(parsed)
470 }
471
472 pub fn matches(&self, scoping: ItemScoping) -> bool {
477 self.matches_scope(scoping)
478 && scoping.matches_categories(&self.categories)
479 && scoping.matches_namespaces(&self.namespace)
480 }
481}
482
483#[cfg(test)]
484mod tests {
485 use smallvec::smallvec;
486
487 use super::*;
488
489 #[test]
490 fn test_parse_quota_reject_all() {
491 let json = r#"{
492 "limit": 0,
493 "reasonCode": "not_yet"
494 }"#;
495
496 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
497
498 insta::assert_ron_snapshot!(quota, @r###"
499 Quota(
500 id: None,
501 categories: [],
502 scope: organization,
503 limit: Some(0),
504 namespace: None,
505 reasonCode: Some(ReasonCode("not_yet")),
506 )
507 "###);
508 }
509
510 #[test]
511 fn test_parse_quota_reject_transactions() {
512 let json = r#"{
513 "limit": 0,
514 "categories": ["transaction"],
515 "reasonCode": "not_yet"
516 }"#;
517
518 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
519
520 insta::assert_ron_snapshot!(quota, @r###"
521 Quota(
522 id: None,
523 categories: [
524 transaction,
525 ],
526 scope: organization,
527 limit: Some(0),
528 namespace: None,
529 reasonCode: Some(ReasonCode("not_yet")),
530 )
531 "###);
532 }
533
534 #[test]
535 fn test_parse_quota_limited() {
536 let json = r#"{
537 "id": "o",
538 "limit": 4711,
539 "window": 42,
540 "reasonCode": "not_so_fast"
541 }"#;
542
543 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
544
545 insta::assert_ron_snapshot!(quota, @r###"
546 Quota(
547 id: Some("o"),
548 categories: [],
549 scope: organization,
550 limit: Some(4711),
551 window: Some(42),
552 namespace: None,
553 reasonCode: Some(ReasonCode("not_so_fast")),
554 )
555 "###);
556 }
557
558 #[test]
559 fn test_parse_quota_project() {
560 let json = r#"{
561 "id": "p",
562 "scope": "project",
563 "scopeId": "1",
564 "limit": 4711,
565 "window": 42,
566 "reasonCode": "not_so_fast"
567 }"#;
568
569 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
570
571 insta::assert_ron_snapshot!(quota, @r###"
572 Quota(
573 id: Some("p"),
574 categories: [],
575 scope: project,
576 scopeId: Some("1"),
577 limit: Some(4711),
578 window: Some(42),
579 namespace: None,
580 reasonCode: Some(ReasonCode("not_so_fast")),
581 )
582 "###);
583 }
584
585 #[test]
586 fn test_parse_quota_project_large() {
587 let json = r#"{
588 "id": "p",
589 "scope": "project",
590 "scopeId": "1",
591 "limit": 4294967296,
592 "window": 42,
593 "reasonCode": "not_so_fast"
594 }"#;
595
596 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
597
598 insta::assert_ron_snapshot!(quota, @r###"
599 Quota(
600 id: Some("p"),
601 categories: [],
602 scope: project,
603 scopeId: Some("1"),
604 limit: Some(4294967296),
605 window: Some(42),
606 namespace: None,
607 reasonCode: Some(ReasonCode("not_so_fast")),
608 )
609 "###);
610 }
611
612 #[test]
613 fn test_parse_quota_key() {
614 let json = r#"{
615 "id": "k",
616 "scope": "key",
617 "scopeId": "1",
618 "limit": 4711,
619 "window": 42,
620 "reasonCode": "not_so_fast"
621 }"#;
622
623 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
624
625 insta::assert_ron_snapshot!(quota, @r###"
626 Quota(
627 id: Some("k"),
628 categories: [],
629 scope: key,
630 scopeId: Some("1"),
631 limit: Some(4711),
632 window: Some(42),
633 namespace: None,
634 reasonCode: Some(ReasonCode("not_so_fast")),
635 )
636 "###);
637 }
638
639 #[test]
640 fn test_parse_quota_unknown_variants() {
641 let json = r#"{
642 "id": "f",
643 "categories": ["future"],
644 "scope": "future",
645 "scopeId": "1",
646 "limit": 4711,
647 "window": 42,
648 "reasonCode": "not_so_fast"
649 }"#;
650
651 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
652
653 insta::assert_ron_snapshot!(quota, @r###"
654 Quota(
655 id: Some("f"),
656 categories: [
657 unknown,
658 ],
659 scope: unknown,
660 scopeId: Some("1"),
661 limit: Some(4711),
662 window: Some(42),
663 namespace: None,
664 reasonCode: Some(ReasonCode("not_so_fast")),
665 )
666 "###);
667 }
668
669 #[test]
670 fn test_parse_quota_unlimited() {
671 let json = r#"{
672 "id": "o",
673 "window": 42
674 }"#;
675
676 let quota = serde_json::from_str::<Quota>(json).expect("parse quota");
677
678 insta::assert_ron_snapshot!(quota, @r###"
679 Quota(
680 id: Some("o"),
681 categories: [],
682 scope: organization,
683 limit: None,
684 window: Some(42),
685 namespace: None,
686 )
687 "###);
688 }
689
690 #[test]
691 fn test_quota_valid_reject_all() {
692 let quota = Quota {
693 id: None,
694 categories: DataCategories::new(),
695 scope: QuotaScope::Organization,
696 scope_id: None,
697 limit: Some(0),
698 window: None,
699 reason_code: None,
700 namespace: None,
701 };
702
703 assert!(quota.is_valid());
704 }
705
706 #[test]
707 fn test_quota_invalid_only_unknown() {
708 let quota = Quota {
709 id: None,
710 categories: smallvec![DataCategory::Unknown, DataCategory::Unknown],
711 scope: QuotaScope::Organization,
712 scope_id: None,
713 limit: Some(0),
714 window: None,
715 reason_code: None,
716 namespace: None,
717 };
718
719 assert!(!quota.is_valid());
720 }
721
722 #[test]
723 fn test_quota_valid_reject_all_mixed() {
724 let quota = Quota {
725 id: None,
726 categories: smallvec![DataCategory::Error, DataCategory::Attachment],
727 scope: QuotaScope::Organization,
728 scope_id: None,
729 limit: Some(0),
730 window: None,
731 reason_code: None,
732 namespace: None,
733 };
734
735 assert!(quota.is_valid());
736 }
737
738 #[test]
739 fn test_quota_invalid_limited_mixed() {
740 let quota = Quota {
741 id: None,
742 categories: smallvec![DataCategory::Error, DataCategory::Attachment],
743 scope: QuotaScope::Organization,
744 scope_id: None,
745 limit: Some(1000),
746 window: None,
747 reason_code: None,
748 namespace: None,
749 };
750
751 assert!(!quota.is_valid());
753 }
754
755 #[test]
756 fn test_quota_invalid_unlimited_mixed() {
757 let quota = Quota {
758 id: None,
759 categories: smallvec![DataCategory::Error, DataCategory::Attachment],
760 scope: QuotaScope::Organization,
761 scope_id: None,
762 limit: None,
763 window: None,
764 reason_code: None,
765 namespace: None,
766 };
767
768 assert!(!quota.is_valid());
770 }
771
772 #[test]
773 fn test_quota_matches_no_categories() {
774 let quota = Quota {
775 id: None,
776 categories: DataCategories::new(),
777 scope: QuotaScope::Organization,
778 scope_id: None,
779 limit: None,
780 window: None,
781 reason_code: None,
782 namespace: None,
783 };
784
785 assert!(quota.matches(ItemScoping {
786 category: DataCategory::Error,
787 scoping: Scoping {
788 organization_id: OrganizationId::new(42),
789 project_id: ProjectId::new(21),
790 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
791 key_id: Some(17),
792 },
793 namespace: MetricNamespaceScoping::None,
794 }));
795 }
796
797 #[test]
798 fn test_quota_matches_unknown_category() {
799 let quota = Quota {
800 id: None,
801 categories: smallvec![DataCategory::Unknown],
802 scope: QuotaScope::Organization,
803 scope_id: None,
804 limit: None,
805 window: None,
806 reason_code: None,
807 namespace: None,
808 };
809
810 assert!(!quota.matches(ItemScoping {
811 category: DataCategory::Error,
812 scoping: Scoping {
813 organization_id: OrganizationId::new(42),
814 project_id: ProjectId::new(21),
815 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
816 key_id: Some(17),
817 },
818 namespace: MetricNamespaceScoping::None,
819 }));
820 }
821
822 #[test]
823 fn test_quota_matches_multiple_categores() {
824 let quota = Quota {
825 id: None,
826 categories: smallvec![DataCategory::Unknown, DataCategory::Error],
827 scope: QuotaScope::Organization,
828 scope_id: None,
829 limit: None,
830 window: None,
831 reason_code: None,
832 namespace: None,
833 };
834
835 assert!(quota.matches(ItemScoping {
836 category: DataCategory::Error,
837 scoping: Scoping {
838 organization_id: OrganizationId::new(42),
839 project_id: ProjectId::new(21),
840 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
841 key_id: Some(17),
842 },
843 namespace: MetricNamespaceScoping::None,
844 }));
845
846 assert!(!quota.matches(ItemScoping {
847 category: DataCategory::Transaction,
848 scoping: Scoping {
849 organization_id: OrganizationId::new(42),
850 project_id: ProjectId::new(21),
851 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
852 key_id: Some(17),
853 },
854 namespace: MetricNamespaceScoping::None,
855 }));
856 }
857
858 #[test]
859 fn test_quota_matches_no_invalid_scope() {
860 let quota = Quota {
861 id: None,
862 categories: DataCategories::new(),
863 scope: QuotaScope::Organization,
864 scope_id: Some("not_a_number".to_owned()),
865 limit: None,
866 window: None,
867 reason_code: None,
868 namespace: None,
869 };
870
871 assert!(!quota.matches(ItemScoping {
872 category: DataCategory::Error,
873 scoping: Scoping {
874 organization_id: OrganizationId::new(42),
875 project_id: ProjectId::new(21),
876 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
877 key_id: Some(17),
878 },
879 namespace: MetricNamespaceScoping::None,
880 }));
881 }
882
883 #[test]
884 fn test_quota_matches_organization_scope() {
885 let quota = Quota {
886 id: None,
887 categories: DataCategories::new(),
888 scope: QuotaScope::Organization,
889 scope_id: Some("42".to_owned()),
890 limit: None,
891 window: None,
892 reason_code: None,
893 namespace: None,
894 };
895
896 assert!(quota.matches(ItemScoping {
897 category: DataCategory::Error,
898 scoping: Scoping {
899 organization_id: OrganizationId::new(42),
900 project_id: ProjectId::new(21),
901 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
902 key_id: Some(17),
903 },
904 namespace: MetricNamespaceScoping::None,
905 }));
906
907 assert!(!quota.matches(ItemScoping {
908 category: DataCategory::Error,
909 scoping: Scoping {
910 organization_id: OrganizationId::new(0),
911 project_id: ProjectId::new(21),
912 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
913 key_id: Some(17),
914 },
915 namespace: MetricNamespaceScoping::None,
916 }));
917 }
918
919 #[test]
920 fn test_quota_matches_project_scope() {
921 let quota = Quota {
922 id: None,
923 categories: DataCategories::new(),
924 scope: QuotaScope::Project,
925 scope_id: Some("21".to_owned()),
926 limit: None,
927 window: None,
928 reason_code: None,
929 namespace: None,
930 };
931
932 assert!(quota.matches(ItemScoping {
933 category: DataCategory::Error,
934 scoping: Scoping {
935 organization_id: OrganizationId::new(42),
936 project_id: ProjectId::new(21),
937 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
938 key_id: Some(17),
939 },
940 namespace: MetricNamespaceScoping::None,
941 }));
942
943 assert!(!quota.matches(ItemScoping {
944 category: DataCategory::Error,
945 scoping: Scoping {
946 organization_id: OrganizationId::new(42),
947 project_id: ProjectId::new(0),
948 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
949 key_id: Some(17),
950 },
951 namespace: MetricNamespaceScoping::None,
952 }));
953 }
954
955 #[test]
956 fn test_quota_matches_key_scope() {
957 let quota = Quota {
958 id: None,
959 categories: DataCategories::new(),
960 scope: QuotaScope::Key,
961 scope_id: Some("17".to_owned()),
962 limit: None,
963 window: None,
964 reason_code: None,
965 namespace: None,
966 };
967
968 assert!(quota.matches(ItemScoping {
969 category: DataCategory::Error,
970 scoping: Scoping {
971 organization_id: OrganizationId::new(42),
972 project_id: ProjectId::new(21),
973 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
974 key_id: Some(17),
975 },
976 namespace: MetricNamespaceScoping::None,
977 }));
978
979 assert!(!quota.matches(ItemScoping {
980 category: DataCategory::Error,
981 scoping: Scoping {
982 organization_id: OrganizationId::new(42),
983 project_id: ProjectId::new(21),
984 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
985 key_id: Some(0),
986 },
987 namespace: MetricNamespaceScoping::None,
988 }));
989
990 assert!(!quota.matches(ItemScoping {
991 category: DataCategory::Error,
992 scoping: Scoping {
993 organization_id: OrganizationId::new(42),
994 project_id: ProjectId::new(21),
995 project_key: ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap(),
996 key_id: None,
997 },
998 namespace: MetricNamespaceScoping::None,
999 }));
1000 }
1001}