nautilus_persistence/writer/
filter.rs1use ahash::{AHashMap, AHashSet};
19use nautilus_model::{
20 data::{NautilusDataType, NautilusRecordType},
21 instruments::NautilusInstrumentType,
22};
23
24use crate::catalog::types::CatalogDataType;
25
26#[derive(Clone, Debug, Default, PartialEq)]
28pub struct WriterRecordFilter {
29 entries: AHashMap<CatalogDataType, Option<AHashSet<String>>>,
30 instrument_types: AHashSet<NautilusInstrumentType>,
31}
32
33impl WriterRecordFilter {
34 #[must_use]
36 pub fn new() -> Self {
37 Self::default()
38 }
39
40 #[must_use]
42 pub fn from_record_types(record_types: impl IntoIterator<Item = NautilusRecordType>) -> Self {
43 let mut filter = Self::new();
44 for record_type in record_types {
45 filter.insert(record_type, None);
46 }
47
48 filter
49 }
50
51 pub fn insert(&mut self, family: impl Into<CatalogDataType>, identifiers: Option<Vec<String>>) {
53 let identifiers = identifiers.map(|values| values.into_iter().collect());
54 self.entries.insert(family.into(), identifiers);
55 }
56
57 pub fn insert_instrument_type(&mut self, instrument_type: NautilusInstrumentType) {
59 self.instrument_types.insert(instrument_type);
60 }
61
62 #[must_use]
64 pub fn is_empty(&self) -> bool {
65 self.entries.is_empty() && self.instrument_types.is_empty()
66 }
67
68 #[must_use]
70 pub fn contains(&self, family: &CatalogDataType) -> bool {
71 self.is_empty()
72 || self.entries.contains_key(family)
73 || (!self.instrument_types.is_empty() && is_instrument_family(family))
74 }
75
76 #[must_use]
80 pub fn allows(&self, data_type: &CatalogDataType, identifier: Option<&str>) -> bool {
81 if self.is_empty() {
82 return true;
83 }
84
85 let instrument_type = match data_type {
86 CatalogDataType::Instrument(instrument_type) => Some(instrument_type),
87 _ => None,
88 };
89
90 let family = catalog_family(data_type);
91
92 let Some(identifiers) = self.entries.get(&family) else {
93 return is_instrument_family(&family)
94 && instrument_type.is_some_and(|value| self.instrument_types.contains(value));
95 };
96
97 if is_instrument_family(&family)
98 && !self.instrument_types.is_empty()
99 && !instrument_type.is_some_and(|value| self.instrument_types.contains(value))
100 {
101 return false;
102 }
103
104 match identifiers {
105 None => true,
106 Some(identifiers) => {
107 identifier.is_some_and(|identifier| identifiers.contains(identifier))
108 }
109 }
110 }
111}
112
113#[must_use]
115pub(crate) fn catalog_family(data_type: &CatalogDataType) -> CatalogDataType {
116 match data_type {
117 CatalogDataType::Instrument(_) => CatalogDataType::Data(NautilusDataType::Instrument),
118 other => other.clone(),
119 }
120}
121
122fn is_instrument_family(family: &CatalogDataType) -> bool {
123 *family == CatalogDataType::Data(NautilusDataType::Instrument)
124}