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}