Skip to main content

aegis_document/
query.rs

1//! Aegis Document Query
2//!
3//! Query language for document filtering and retrieval.
4//!
5//! @version 0.1.0
6//! @author AutomataNexus Development Team
7
8use crate::types::{Document, Value};
9use serde::{Deserialize, Serialize};
10
11// =============================================================================
12// Query
13// =============================================================================
14
15/// A query for filtering documents.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Query {
18    pub filters: Vec<Filter>,
19    pub sort: Option<Sort>,
20    pub skip: Option<usize>,
21    pub limit: Option<usize>,
22    pub projection: Option<Vec<String>>,
23}
24
25impl Query {
26    pub fn new() -> Self {
27        Self {
28            filters: Vec::new(),
29            sort: None,
30            skip: None,
31            limit: None,
32            projection: None,
33        }
34    }
35
36    /// Check if a document matches this query.
37    pub fn matches(&self, doc: &Document) -> bool {
38        self.filters.iter().all(|f| f.matches(doc))
39    }
40
41    /// Add a filter.
42    pub fn with_filter(mut self, filter: Filter) -> Self {
43        self.filters.push(filter);
44        self
45    }
46
47    /// Add sorting.
48    pub fn with_sort(mut self, field: impl Into<String>, ascending: bool) -> Self {
49        self.sort = Some(Sort {
50            field: field.into(),
51            ascending,
52        });
53        self
54    }
55
56    /// Add skip.
57    pub fn with_skip(mut self, skip: usize) -> Self {
58        self.skip = Some(skip);
59        self
60    }
61
62    /// Add limit.
63    pub fn with_limit(mut self, limit: usize) -> Self {
64        self.limit = Some(limit);
65        self
66    }
67
68    /// Add projection.
69    pub fn with_projection(mut self, fields: Vec<String>) -> Self {
70        self.projection = Some(fields);
71        self
72    }
73}
74
75impl Default for Query {
76    fn default() -> Self {
77        Self::new()
78    }
79}
80
81// =============================================================================
82// Filter
83// =============================================================================
84
85/// A filter condition.
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub enum Filter {
88    Eq { field: String, value: Value },
89    Ne { field: String, value: Value },
90    Gt { field: String, value: Value },
91    Gte { field: String, value: Value },
92    Lt { field: String, value: Value },
93    Lte { field: String, value: Value },
94    In { field: String, values: Vec<Value> },
95    Nin { field: String, values: Vec<Value> },
96    Exists { field: String, exists: bool },
97    Regex { field: String, pattern: String },
98    Contains { field: String, value: String },
99    StartsWith { field: String, value: String },
100    EndsWith { field: String, value: String },
101    And(Vec<Filter>),
102    Or(Vec<Filter>),
103    Not(Box<Filter>),
104}
105
106impl Filter {
107    /// Check if a document matches this filter.
108    pub fn matches(&self, doc: &Document) -> bool {
109        match self {
110            Self::Eq { field, value } => doc.get(field).map(|v| v == value).unwrap_or(false),
111            Self::Ne { field, value } => doc.get(field).map(|v| v != value).unwrap_or(true),
112            Self::Gt { field, value } => doc
113                .get(field)
114                .map(|v| compare_values(v, value) == Some(std::cmp::Ordering::Greater))
115                .unwrap_or(false),
116            Self::Gte { field, value } => doc
117                .get(field)
118                .map(|v| {
119                    matches!(
120                        compare_values(v, value),
121                        Some(std::cmp::Ordering::Greater | std::cmp::Ordering::Equal)
122                    )
123                })
124                .unwrap_or(false),
125            Self::Lt { field, value } => doc
126                .get(field)
127                .map(|v| compare_values(v, value) == Some(std::cmp::Ordering::Less))
128                .unwrap_or(false),
129            Self::Lte { field, value } => doc
130                .get(field)
131                .map(|v| {
132                    matches!(
133                        compare_values(v, value),
134                        Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal)
135                    )
136                })
137                .unwrap_or(false),
138            Self::In { field, values } => {
139                doc.get(field).map(|v| values.contains(v)).unwrap_or(false)
140            }
141            Self::Nin { field, values } => {
142                doc.get(field).map(|v| !values.contains(v)).unwrap_or(true)
143            }
144            Self::Exists { field, exists } => doc.contains(field) == *exists,
145            Self::Regex { field, pattern } => {
146                if let Some(Value::String(s)) = doc.get(field) {
147                    regex::Regex::new(pattern)
148                        .map(|re| re.is_match(s))
149                        .unwrap_or(false)
150                } else {
151                    false
152                }
153            }
154            Self::Contains { field, value } => {
155                if let Some(Value::String(s)) = doc.get(field) {
156                    s.contains(value)
157                } else {
158                    false
159                }
160            }
161            Self::StartsWith { field, value } => {
162                if let Some(Value::String(s)) = doc.get(field) {
163                    s.starts_with(value)
164                } else {
165                    false
166                }
167            }
168            Self::EndsWith { field, value } => {
169                if let Some(Value::String(s)) = doc.get(field) {
170                    s.ends_with(value)
171                } else {
172                    false
173                }
174            }
175            Self::And(filters) => filters.iter().all(|f| f.matches(doc)),
176            Self::Or(filters) => filters.iter().any(|f| f.matches(doc)),
177            Self::Not(filter) => !filter.matches(doc),
178        }
179    }
180}
181
182fn compare_values(a: &Value, b: &Value) -> Option<std::cmp::Ordering> {
183    match (a, b) {
184        (Value::Int(a), Value::Int(b)) => Some(a.cmp(b)),
185        (Value::Float(a), Value::Float(b)) => a.partial_cmp(b),
186        (Value::Int(a), Value::Float(b)) => (*a as f64).partial_cmp(b),
187        (Value::Float(a), Value::Int(b)) => a.partial_cmp(&(*b as f64)),
188        (Value::String(a), Value::String(b)) => Some(a.cmp(b)),
189        (Value::Bool(a), Value::Bool(b)) => Some(a.cmp(b)),
190        _ => None,
191    }
192}
193
194// =============================================================================
195// Sort
196// =============================================================================
197
198/// Sort specification.
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct Sort {
201    pub field: String,
202    pub ascending: bool,
203}
204
205// =============================================================================
206// Query Builder
207// =============================================================================
208
209/// Builder for constructing queries.
210pub struct QueryBuilder {
211    query: Query,
212}
213
214impl QueryBuilder {
215    pub fn new() -> Self {
216        Self {
217            query: Query::new(),
218        }
219    }
220
221    pub fn eq(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
222        self.query.filters.push(Filter::Eq {
223            field: field.into(),
224            value: value.into(),
225        });
226        self
227    }
228
229    pub fn ne(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
230        self.query.filters.push(Filter::Ne {
231            field: field.into(),
232            value: value.into(),
233        });
234        self
235    }
236
237    pub fn gt(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
238        self.query.filters.push(Filter::Gt {
239            field: field.into(),
240            value: value.into(),
241        });
242        self
243    }
244
245    pub fn gte(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
246        self.query.filters.push(Filter::Gte {
247            field: field.into(),
248            value: value.into(),
249        });
250        self
251    }
252
253    pub fn lt(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
254        self.query.filters.push(Filter::Lt {
255            field: field.into(),
256            value: value.into(),
257        });
258        self
259    }
260
261    pub fn lte(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
262        self.query.filters.push(Filter::Lte {
263            field: field.into(),
264            value: value.into(),
265        });
266        self
267    }
268
269    pub fn in_values(mut self, field: impl Into<String>, values: Vec<Value>) -> Self {
270        self.query.filters.push(Filter::In {
271            field: field.into(),
272            values,
273        });
274        self
275    }
276
277    pub fn exists(mut self, field: impl Into<String>, exists: bool) -> Self {
278        self.query.filters.push(Filter::Exists {
279            field: field.into(),
280            exists,
281        });
282        self
283    }
284
285    pub fn contains(mut self, field: impl Into<String>, value: impl Into<String>) -> Self {
286        self.query.filters.push(Filter::Contains {
287            field: field.into(),
288            value: value.into(),
289        });
290        self
291    }
292
293    pub fn starts_with(mut self, field: impl Into<String>, value: impl Into<String>) -> Self {
294        self.query.filters.push(Filter::StartsWith {
295            field: field.into(),
296            value: value.into(),
297        });
298        self
299    }
300
301    pub fn ends_with(mut self, field: impl Into<String>, value: impl Into<String>) -> Self {
302        self.query.filters.push(Filter::EndsWith {
303            field: field.into(),
304            value: value.into(),
305        });
306        self
307    }
308
309    pub fn regex(mut self, field: impl Into<String>, pattern: impl Into<String>) -> Self {
310        self.query.filters.push(Filter::Regex {
311            field: field.into(),
312            pattern: pattern.into(),
313        });
314        self
315    }
316
317    pub fn and(mut self, filters: Vec<Filter>) -> Self {
318        self.query.filters.push(Filter::And(filters));
319        self
320    }
321
322    pub fn or(mut self, filters: Vec<Filter>) -> Self {
323        self.query.filters.push(Filter::Or(filters));
324        self
325    }
326
327    pub fn sort(mut self, field: impl Into<String>, ascending: bool) -> Self {
328        self.query.sort = Some(Sort {
329            field: field.into(),
330            ascending,
331        });
332        self
333    }
334
335    pub fn skip(mut self, skip: usize) -> Self {
336        self.query.skip = Some(skip);
337        self
338    }
339
340    pub fn limit(mut self, limit: usize) -> Self {
341        self.query.limit = Some(limit);
342        self
343    }
344
345    pub fn project(mut self, fields: Vec<String>) -> Self {
346        self.query.projection = Some(fields);
347        self
348    }
349
350    pub fn build(self) -> Query {
351        self.query
352    }
353}
354
355impl Default for QueryBuilder {
356    fn default() -> Self {
357        Self::new()
358    }
359}
360
361// =============================================================================
362// Query Result
363// =============================================================================
364
365/// Result of a document query.
366#[derive(Debug, Clone)]
367pub struct QueryResult {
368    pub documents: Vec<Document>,
369    pub total_scanned: usize,
370    pub execution_time_ms: u64,
371}
372
373impl QueryResult {
374    pub fn empty() -> Self {
375        Self {
376            documents: Vec::new(),
377            total_scanned: 0,
378            execution_time_ms: 0,
379        }
380    }
381
382    pub fn count(&self) -> usize {
383        self.documents.len()
384    }
385
386    pub fn is_empty(&self) -> bool {
387        self.documents.is_empty()
388    }
389
390    pub fn first(&self) -> Option<&Document> {
391        self.documents.first()
392    }
393}
394
395// =============================================================================
396// Tests
397// =============================================================================
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402
403    fn create_test_doc() -> Document {
404        let mut doc = Document::with_id("test");
405        doc.set("name", "Alice");
406        doc.set("age", 30i64);
407        doc.set("active", true);
408        doc.set("email", "alice@example.com");
409        doc
410    }
411
412    #[test]
413    fn test_eq_filter() {
414        let doc = create_test_doc();
415
416        let filter = Filter::Eq {
417            field: "name".to_string(),
418            value: Value::String("Alice".to_string()),
419        };
420        assert!(filter.matches(&doc));
421
422        let filter = Filter::Eq {
423            field: "name".to_string(),
424            value: Value::String("Bob".to_string()),
425        };
426        assert!(!filter.matches(&doc));
427    }
428
429    #[test]
430    fn test_comparison_filters() {
431        let doc = create_test_doc();
432
433        let filter = Filter::Gt {
434            field: "age".to_string(),
435            value: Value::Int(25),
436        };
437        assert!(filter.matches(&doc));
438
439        let filter = Filter::Lt {
440            field: "age".to_string(),
441            value: Value::Int(25),
442        };
443        assert!(!filter.matches(&doc));
444
445        let filter = Filter::Gte {
446            field: "age".to_string(),
447            value: Value::Int(30),
448        };
449        assert!(filter.matches(&doc));
450    }
451
452    #[test]
453    fn test_string_filters() {
454        let doc = create_test_doc();
455
456        let filter = Filter::Contains {
457            field: "email".to_string(),
458            value: "example".to_string(),
459        };
460        assert!(filter.matches(&doc));
461
462        let filter = Filter::StartsWith {
463            field: "email".to_string(),
464            value: "alice".to_string(),
465        };
466        assert!(filter.matches(&doc));
467
468        let filter = Filter::EndsWith {
469            field: "email".to_string(),
470            value: ".com".to_string(),
471        };
472        assert!(filter.matches(&doc));
473    }
474
475    #[test]
476    fn test_logical_filters() {
477        let doc = create_test_doc();
478
479        let filter = Filter::And(vec![
480            Filter::Eq {
481                field: "name".to_string(),
482                value: Value::String("Alice".to_string()),
483            },
484            Filter::Gt {
485                field: "age".to_string(),
486                value: Value::Int(20),
487            },
488        ]);
489        assert!(filter.matches(&doc));
490
491        let filter = Filter::Or(vec![
492            Filter::Eq {
493                field: "name".to_string(),
494                value: Value::String("Bob".to_string()),
495            },
496            Filter::Eq {
497                field: "active".to_string(),
498                value: Value::Bool(true),
499            },
500        ]);
501        assert!(filter.matches(&doc));
502    }
503
504    #[test]
505    fn test_query_builder() {
506        let doc = create_test_doc();
507
508        let query = QueryBuilder::new()
509            .eq("name", "Alice")
510            .gt("age", 25i64)
511            .build();
512
513        assert!(query.matches(&doc));
514
515        let query = QueryBuilder::new().eq("name", "Bob").build();
516
517        assert!(!query.matches(&doc));
518    }
519
520    #[test]
521    fn test_exists_filter() {
522        let doc = create_test_doc();
523
524        let filter = Filter::Exists {
525            field: "name".to_string(),
526            exists: true,
527        };
528        assert!(filter.matches(&doc));
529
530        let filter = Filter::Exists {
531            field: "missing".to_string(),
532            exists: false,
533        };
534        assert!(filter.matches(&doc));
535    }
536}