exocore_store/
query.rs

1use exocore_protos::{
2    generated::exocore_store::{
3        entity_query, ordering, trait_field_predicate, trait_query, EntityQuery, MatchPredicate,
4        Ordering, Paging, ReferencePredicate, TraitFieldPredicate, TraitFieldReferencePredicate,
5        TraitPredicate, TraitQuery,
6    },
7    message::NamedMessage,
8    reflect::FieldId,
9    store::{
10        AllPredicate, IdsPredicate, OperationsPredicate, Projection, QueryStringPredicate,
11        Reference,
12    },
13};
14
15use crate::{
16    entity::{EntityId, TraitId},
17    mutation::OperationId,
18};
19
20pub type WatchToken = u64;
21pub type ResultHash = u64;
22
23#[derive(Clone)]
24pub struct QueryBuilder {
25    query: EntityQuery,
26}
27
28impl QueryBuilder {
29    pub fn matches<T: Into<String>>(query: T) -> QueryBuilder {
30        QueryBuilder {
31            query: EntityQuery {
32                predicate: Some(entity_query::Predicate::Match(MatchPredicate {
33                    query: query.into(),
34                    ..Default::default()
35                })),
36                ..Default::default()
37            },
38        }
39    }
40
41    pub fn references<T: Into<ReferencePredicateWrapper>>(reference: T) -> QueryBuilder {
42        QueryBuilder {
43            query: EntityQuery {
44                predicate: Some(entity_query::Predicate::Reference(reference.into().0)),
45                ..Default::default()
46            },
47        }
48    }
49
50    pub fn with_operations(operation_ids: Vec<OperationId>) -> QueryBuilder {
51        QueryBuilder {
52            query: EntityQuery {
53                predicate: Some(entity_query::Predicate::Operations(OperationsPredicate {
54                    operation_ids,
55                })),
56                ..Default::default()
57            },
58        }
59    }
60
61    pub fn with_trait_name<T: Into<String>>(trait_name: T) -> QueryBuilder {
62        QueryBuilder {
63            query: EntityQuery {
64                predicate: Some(entity_query::Predicate::Trait(TraitPredicate {
65                    trait_name: trait_name.into(),
66                    ..Default::default()
67                })),
68                ..Default::default()
69            },
70        }
71    }
72
73    pub fn with_trait<T: NamedMessage>() -> QueryBuilder {
74        Self::with_trait_name(T::full_name())
75    }
76
77    pub fn with_trait_name_query<T: Into<String>>(
78        trait_name: T,
79        query: TraitQuery,
80    ) -> QueryBuilder {
81        QueryBuilder {
82            query: EntityQuery {
83                predicate: Some(entity_query::Predicate::Trait(TraitPredicate {
84                    trait_name: trait_name.into(),
85                    query: Some(query),
86                })),
87                ..Default::default()
88            },
89        }
90    }
91
92    pub fn with_trait_query<T: NamedMessage>(query: TraitQuery) -> QueryBuilder {
93        Self::with_trait_name_query(T::full_name(), query)
94    }
95
96    pub fn with_id<E: Into<String>>(id: E) -> QueryBuilder {
97        QueryBuilder {
98            query: EntityQuery {
99                predicate: Some(entity_query::Predicate::Ids(IdsPredicate {
100                    ids: vec![id.into()],
101                })),
102                ..Default::default()
103            },
104        }
105    }
106
107    pub fn with_ids<I>(ids: I) -> QueryBuilder
108    where
109        I: IntoIterator,
110        I::Item: Into<String>,
111    {
112        QueryBuilder {
113            query: EntityQuery {
114                predicate: Some(entity_query::Predicate::Ids(IdsPredicate {
115                    ids: ids.into_iter().map(|i| i.into()).collect(),
116                })),
117                ..Default::default()
118            },
119        }
120    }
121
122    pub fn from_query_string<T: Into<String>>(query: T) -> QueryBuilder {
123        QueryBuilder {
124            query: EntityQuery {
125                predicate: Some(entity_query::Predicate::QueryString(QueryStringPredicate {
126                    query: query.into(),
127                })),
128                ..Default::default()
129            },
130        }
131    }
132
133    pub fn all() -> QueryBuilder {
134        QueryBuilder {
135            query: EntityQuery {
136                predicate: Some(entity_query::Predicate::All(AllPredicate {})),
137                ..Default::default()
138            },
139        }
140    }
141
142    #[cfg(any(test, feature = "tests-utils"))]
143    pub fn test(success: bool) -> QueryBuilder {
144        use exocore_protos::generated::exocore_store::TestPredicate;
145        QueryBuilder {
146            query: EntityQuery {
147                predicate: Some(entity_query::Predicate::Test(TestPredicate { success })),
148                ..Default::default()
149            },
150        }
151    }
152
153    pub fn with_paging(mut self, paging: Paging) -> Self {
154        self.query.paging = Some(paging);
155        self
156    }
157
158    pub fn count(mut self, count: u32) -> Self {
159        match self.query.paging.as_mut() {
160            Some(paging) => paging.count = count,
161            None => {
162                self.query.paging = Some(Paging {
163                    count,
164                    ..Default::default()
165                })
166            }
167        }
168
169        self
170    }
171
172    pub fn project<P: Into<ProjectionWrapper>>(mut self, projection: P) -> Self {
173        self.query.projections.push(projection.into().0);
174        self
175    }
176
177    pub fn projects<I>(mut self, projections: I) -> Self
178    where
179        I: IntoIterator,
180        I::Item: Into<ProjectionWrapper>,
181    {
182        self.query.projections = projections.into_iter().map(|p| p.into().0).collect();
183        self
184    }
185
186    pub fn skip_if_results_equals(mut self, result_hash: ResultHash) -> Self {
187        self.query.result_hash = result_hash;
188        self
189    }
190
191    pub fn paging_or_default(&self) -> Paging {
192        self.query.paging.clone().unwrap_or_else(default_paging)
193    }
194
195    pub fn with_watch_token(mut self, token: WatchToken) -> Self {
196        self.query.watch_token = token;
197        self
198    }
199
200    pub fn order_by_field<F: Into<String>>(self, field: F, ascending: bool) -> Self {
201        self.mapped_ordering(|ordering| {
202            ordering.value = Some(ordering::Value::Field(field.into()));
203            ordering.ascending = ascending;
204        })
205    }
206
207    pub fn order_by_operations(self, ascending: bool) -> Self {
208        self.mapped_ordering(|ordering| {
209            ordering.value = Some(ordering::Value::OperationId(true));
210            ordering.ascending = ascending;
211        })
212    }
213
214    pub fn order_by_score(
215        self,
216        ascending: bool,
217        recency_boost: bool,
218        reference_boost: bool,
219    ) -> Self {
220        self.mapped_ordering(|ordering| {
221            ordering.value = Some(ordering::Value::Score(true));
222            ordering.ascending = ascending;
223            ordering.no_recency_boost = !recency_boost;
224            ordering.no_reference_boost = !reference_boost;
225        })
226    }
227
228    pub fn order_ascending(self, ascending: bool) -> Self {
229        self.mapped_ordering(|ordering| ordering.ascending = ascending)
230    }
231
232    pub fn mapped_ordering<F: FnOnce(&mut Ordering)>(mut self, f: F) -> Self {
233        match self.query.ordering.as_mut() {
234            Some(ordering) => f(ordering),
235            None => {
236                let mut ordering = Ordering::default();
237                f(&mut ordering);
238                self.query.ordering = Some(ordering);
239            }
240        }
241
242        self
243    }
244
245    pub fn include_deleted(mut self) -> Self {
246        self.query.include_deleted = true;
247        self
248    }
249
250    pub fn programmatic(mut self) -> Self {
251        self.query.programmatic = true;
252        self
253    }
254
255    pub fn build(self) -> EntityQuery {
256        self.query
257    }
258}
259
260pub struct TraitQueryBuilder {
261    query: TraitQuery,
262}
263
264impl TraitQueryBuilder {
265    pub fn matches<S: Into<String>>(query: S) -> TraitQueryBuilder {
266        TraitQueryBuilder {
267            query: TraitQuery {
268                predicate: Some(trait_query::Predicate::Match(MatchPredicate {
269                    query: query.into(),
270                    ..Default::default()
271                })),
272            },
273        }
274    }
275
276    pub fn field_equals<F: Into<String>, V: Into<FieldPredicateValueWrapper>>(
277        field: F,
278        value: V,
279    ) -> TraitQueryBuilder {
280        TraitQueryBuilder {
281            query: TraitQuery {
282                predicate: Some(trait_query::Predicate::Field(TraitFieldPredicate {
283                    field: field.into(),
284                    value: Some(value.into().0),
285                    operator: trait_field_predicate::Operator::Equal.into(),
286                })),
287            },
288        }
289    }
290
291    pub fn field_references<F: Into<String>, V: Into<ReferencePredicateWrapper>>(
292        field: F,
293        reference: V,
294    ) -> TraitQueryBuilder {
295        TraitQueryBuilder {
296            query: TraitQuery {
297                predicate: Some(trait_query::Predicate::Reference(
298                    TraitFieldReferencePredicate {
299                        field: field.into(),
300                        reference: Some(reference.into().0),
301                    },
302                )),
303            },
304        }
305    }
306
307    pub fn from_query_string<S: Into<String>>(query: S) -> TraitQueryBuilder {
308        TraitQueryBuilder {
309            query: TraitQuery {
310                predicate: Some(trait_query::Predicate::QueryString(QueryStringPredicate {
311                    query: query.into(),
312                })),
313            },
314        }
315    }
316
317    pub fn build(self) -> TraitQuery {
318        self.query
319    }
320}
321
322pub struct ProjectionBuilder {
323    projection: Projection,
324}
325
326impl ProjectionBuilder {
327    pub fn for_package_prefix<S: Into<String>>(package: S) -> ProjectionBuilder {
328        ProjectionBuilder {
329            projection: Projection {
330                package: vec![package.into()],
331                ..Default::default()
332            },
333        }
334    }
335
336    pub fn for_trait_name<S: Into<String>>(trait_name: S) -> ProjectionBuilder {
337        ProjectionBuilder {
338            projection: Projection {
339                package: vec![format!("{}$", trait_name.into())],
340                ..Default::default()
341            },
342        }
343    }
344
345    pub fn for_trait<T: NamedMessage>() -> ProjectionBuilder {
346        Self::for_trait_name(T::full_name())
347    }
348
349    pub fn for_all() -> ProjectionBuilder {
350        ProjectionBuilder {
351            projection: Default::default(),
352        }
353    }
354
355    pub fn skip(mut self) -> ProjectionBuilder {
356        self.projection.skip = true;
357        self
358    }
359
360    pub fn return_all(self) -> ProjectionBuilder {
361        self
362    }
363
364    pub fn return_fields(mut self, field_ids: Vec<FieldId>) -> ProjectionBuilder {
365        self.projection.field_ids = field_ids;
366        self
367    }
368
369    pub fn return_field_groups(mut self, field_groups: Vec<FieldId>) -> ProjectionBuilder {
370        self.projection.field_group_ids = field_groups;
371        self
372    }
373
374    pub fn build(self) -> Projection {
375        self.projection
376    }
377}
378
379pub struct FieldPredicateValueWrapper(trait_field_predicate::Value);
380
381impl From<EntityId> for FieldPredicateValueWrapper {
382    fn from(id: EntityId) -> Self {
383        FieldPredicateValueWrapper(trait_field_predicate::Value::String(id))
384    }
385}
386
387impl From<&str> for FieldPredicateValueWrapper {
388    fn from(s: &str) -> Self {
389        FieldPredicateValueWrapper(trait_field_predicate::Value::String(s.to_string()))
390    }
391}
392
393pub struct ReferencePredicateWrapper(ReferencePredicate);
394
395impl From<EntityId> for ReferencePredicateWrapper {
396    fn from(id: EntityId) -> Self {
397        ReferencePredicateWrapper(ReferencePredicate {
398            entity_id: id,
399            trait_id: String::new(),
400        })
401    }
402}
403
404impl From<(EntityId, TraitId)> for ReferencePredicateWrapper {
405    fn from(tup: (EntityId, TraitId)) -> Self {
406        ReferencePredicateWrapper(ReferencePredicate {
407            entity_id: tup.0,
408            trait_id: tup.1,
409        })
410    }
411}
412
413impl From<&str> for ReferencePredicateWrapper {
414    fn from(s: &str) -> Self {
415        ReferencePredicateWrapper(ReferencePredicate {
416            entity_id: s.to_string(),
417            trait_id: String::new(),
418        })
419    }
420}
421
422impl From<(&str, &str)> for ReferencePredicateWrapper {
423    fn from(tup: (&str, &str)) -> Self {
424        ReferencePredicateWrapper(ReferencePredicate {
425            entity_id: tup.0.to_string(),
426            trait_id: tup.1.to_string(),
427        })
428    }
429}
430
431impl From<Reference> for ReferencePredicateWrapper {
432    fn from(r: Reference) -> Self {
433        ReferencePredicateWrapper(ReferencePredicate {
434            entity_id: r.entity_id,
435            trait_id: r.trait_id,
436        })
437    }
438}
439
440pub struct ProjectionWrapper(Projection);
441
442impl From<Projection> for ProjectionWrapper {
443    fn from(p: Projection) -> Self {
444        ProjectionWrapper(p)
445    }
446}
447
448impl From<ProjectionBuilder> for ProjectionWrapper {
449    fn from(b: ProjectionBuilder) -> Self {
450        ProjectionWrapper(b.build())
451    }
452}
453
454pub fn default_paging() -> Paging {
455    Paging {
456        count: 10,
457        ..Default::default()
458    }
459}
460
461pub fn fill_default_paging(paging: &mut Paging) {
462    if paging.count == 0 {
463        paging.count = 10;
464    }
465}