aegis_document/
collection.rs

1//! Aegis Document Collection
2//!
3//! Collection management for document storage.
4//!
5//! @version 0.1.0
6//! @author AutomataNexus Development Team
7
8use crate::index::{DocumentIndex, IndexType};
9use crate::query::{Query, QueryResult};
10use crate::types::{Document, DocumentId};
11use crate::validation::{Schema, ValidationResult};
12use std::collections::HashMap;
13use std::sync::RwLock;
14
15// =============================================================================
16// Collection
17// =============================================================================
18
19/// A collection of documents.
20pub struct Collection {
21    name: String,
22    documents: RwLock<HashMap<DocumentId, Document>>,
23    indexes: RwLock<Vec<DocumentIndex>>,
24    schema: Option<Schema>,
25}
26
27impl Collection {
28    /// Create a new collection.
29    pub fn new(name: impl Into<String>) -> Self {
30        Self {
31            name: name.into(),
32            documents: RwLock::new(HashMap::new()),
33            indexes: RwLock::new(Vec::new()),
34            schema: None,
35        }
36    }
37
38    /// Create a collection with schema validation.
39    pub fn with_schema(name: impl Into<String>, schema: Schema) -> Self {
40        Self {
41            name: name.into(),
42            documents: RwLock::new(HashMap::new()),
43            indexes: RwLock::new(Vec::new()),
44            schema: Some(schema),
45        }
46    }
47
48    /// Get the collection name.
49    pub fn name(&self) -> &str {
50        &self.name
51    }
52
53    // -------------------------------------------------------------------------
54    // Document Operations
55    // -------------------------------------------------------------------------
56
57    /// Insert a document.
58    pub fn insert(&self, doc: Document) -> Result<DocumentId, CollectionError> {
59        if let Some(ref schema) = self.schema {
60            let result = schema.validate(&doc);
61            if !result.is_valid {
62                return Err(CollectionError::ValidationFailed(result.errors));
63            }
64        }
65
66        let id = doc.id.clone();
67
68        {
69            let mut docs = self.documents.write().unwrap();
70            if docs.contains_key(&id) {
71                return Err(CollectionError::DuplicateId(id));
72            }
73            docs.insert(id.clone(), doc.clone());
74        }
75
76        self.index_document(&doc);
77
78        Ok(id)
79    }
80
81    /// Insert multiple documents.
82    pub fn insert_many(&self, docs: Vec<Document>) -> Result<Vec<DocumentId>, CollectionError> {
83        let mut ids = Vec::with_capacity(docs.len());
84
85        for doc in docs {
86            let id = self.insert(doc)?;
87            ids.push(id);
88        }
89
90        Ok(ids)
91    }
92
93    /// Get a document by ID.
94    pub fn get(&self, id: &DocumentId) -> Option<Document> {
95        let docs = self.documents.read().unwrap();
96        docs.get(id).cloned()
97    }
98
99    /// Update a document.
100    pub fn update(&self, id: &DocumentId, doc: Document) -> Result<(), CollectionError> {
101        if let Some(ref schema) = self.schema {
102            let result = schema.validate(&doc);
103            if !result.is_valid {
104                return Err(CollectionError::ValidationFailed(result.errors));
105            }
106        }
107
108        {
109            let mut docs = self.documents.write().unwrap();
110            if !docs.contains_key(id) {
111                return Err(CollectionError::NotFound(id.clone()));
112            }
113
114            if let Some(old_doc) = docs.get(id) {
115                self.unindex_document(old_doc);
116            }
117
118            docs.insert(id.clone(), doc.clone());
119        }
120
121        self.index_document(&doc);
122
123        Ok(())
124    }
125
126    /// Delete a document.
127    pub fn delete(&self, id: &DocumentId) -> Result<Document, CollectionError> {
128        let mut docs = self.documents.write().unwrap();
129
130        match docs.remove(id) {
131            Some(doc) => {
132                self.unindex_document(&doc);
133                Ok(doc)
134            }
135            None => Err(CollectionError::NotFound(id.clone())),
136        }
137    }
138
139    /// Check if a document exists.
140    pub fn contains(&self, id: &DocumentId) -> bool {
141        let docs = self.documents.read().unwrap();
142        docs.contains_key(id)
143    }
144
145    /// Get the number of documents.
146    pub fn count(&self) -> usize {
147        let docs = self.documents.read().unwrap();
148        docs.len()
149    }
150
151    /// Get all document IDs.
152    pub fn ids(&self) -> Vec<DocumentId> {
153        let docs = self.documents.read().unwrap();
154        docs.keys().cloned().collect()
155    }
156
157    /// Get all documents.
158    pub fn all(&self) -> Vec<Document> {
159        let docs = self.documents.read().unwrap();
160        docs.values().cloned().collect()
161    }
162
163    /// Clear all documents.
164    pub fn clear(&self) {
165        let mut docs = self.documents.write().unwrap();
166        docs.clear();
167
168        let mut indexes = self.indexes.write().unwrap();
169        for index in indexes.iter_mut() {
170            index.clear();
171        }
172    }
173
174    // -------------------------------------------------------------------------
175    // Query Operations
176    // -------------------------------------------------------------------------
177
178    /// Find documents matching a query.
179    pub fn find(&self, query: &Query) -> QueryResult {
180        let docs = self.documents.read().unwrap();
181        let start = std::time::Instant::now();
182
183        let matching: Vec<Document> = docs
184            .values()
185            .filter(|doc| query.matches(doc))
186            .take(query.limit.unwrap_or(usize::MAX))
187            .cloned()
188            .collect();
189
190        QueryResult {
191            documents: matching,
192            total_scanned: docs.len(),
193            execution_time_ms: start.elapsed().as_millis() as u64,
194        }
195    }
196
197    /// Find one document matching a query.
198    pub fn find_one(&self, query: &Query) -> Option<Document> {
199        let docs = self.documents.read().unwrap();
200        docs.values().find(|doc| query.matches(doc)).cloned()
201    }
202
203    /// Count documents matching a query.
204    pub fn count_matching(&self, query: &Query) -> usize {
205        let docs = self.documents.read().unwrap();
206        docs.values().filter(|doc| query.matches(doc)).count()
207    }
208
209    // -------------------------------------------------------------------------
210    // Index Operations
211    // -------------------------------------------------------------------------
212
213    /// Create an index on a field.
214    pub fn create_index(&self, field: impl Into<String>, index_type: IndexType) {
215        let field = field.into();
216        let mut index = DocumentIndex::new(field.clone(), index_type);
217
218        let docs = self.documents.read().unwrap();
219        for doc in docs.values() {
220            index.index_document(doc);
221        }
222
223        let mut indexes = self.indexes.write().unwrap();
224        indexes.push(index);
225    }
226
227    /// Drop an index.
228    pub fn drop_index(&self, field: &str) {
229        let mut indexes = self.indexes.write().unwrap();
230        indexes.retain(|idx| idx.field() != field);
231    }
232
233    /// Get all index names.
234    pub fn index_names(&self) -> Vec<String> {
235        let indexes = self.indexes.read().unwrap();
236        indexes.iter().map(|idx| idx.field().to_string()).collect()
237    }
238
239    fn index_document(&self, doc: &Document) {
240        let mut indexes = self.indexes.write().unwrap();
241        for index in indexes.iter_mut() {
242            index.index_document(doc);
243        }
244    }
245
246    fn unindex_document(&self, doc: &Document) {
247        let mut indexes = self.indexes.write().unwrap();
248        for index in indexes.iter_mut() {
249            index.unindex_document(doc);
250        }
251    }
252
253    // -------------------------------------------------------------------------
254    // Schema Operations
255    // -------------------------------------------------------------------------
256
257    /// Set the collection schema.
258    pub fn set_schema(&mut self, schema: Schema) {
259        self.schema = Some(schema);
260    }
261
262    /// Get the collection schema.
263    pub fn schema(&self) -> Option<&Schema> {
264        self.schema.as_ref()
265    }
266
267    /// Validate all documents against the schema.
268    pub fn validate_all(&self) -> Vec<(DocumentId, ValidationResult)> {
269        let Some(ref schema) = self.schema else {
270            return Vec::new();
271        };
272
273        let docs = self.documents.read().unwrap();
274        docs.iter()
275            .map(|(id, doc)| (id.clone(), schema.validate(doc)))
276            .filter(|(_, result)| !result.is_valid)
277            .collect()
278    }
279}
280
281// =============================================================================
282// Collection Error
283// =============================================================================
284
285/// Errors that can occur in collection operations.
286#[derive(Debug, Clone)]
287pub enum CollectionError {
288    DuplicateId(DocumentId),
289    NotFound(DocumentId),
290    ValidationFailed(Vec<String>),
291    IndexError(String),
292}
293
294impl std::fmt::Display for CollectionError {
295    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296        match self {
297            Self::DuplicateId(id) => write!(f, "Document with ID {} already exists", id),
298            Self::NotFound(id) => write!(f, "Document with ID {} not found", id),
299            Self::ValidationFailed(errors) => {
300                write!(f, "Validation failed: {}", errors.join(", "))
301            }
302            Self::IndexError(msg) => write!(f, "Index error: {}", msg),
303        }
304    }
305}
306
307impl std::error::Error for CollectionError {}
308
309// =============================================================================
310// Tests
311// =============================================================================
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::query::QueryBuilder;
317
318    #[test]
319    fn test_collection_creation() {
320        let collection = Collection::new("users");
321        assert_eq!(collection.name(), "users");
322        assert_eq!(collection.count(), 0);
323    }
324
325    #[test]
326    fn test_insert_and_get() {
327        let collection = Collection::new("test");
328
329        let mut doc = Document::with_id("doc1");
330        doc.set("name", "Alice");
331
332        let id = collection.insert(doc).unwrap();
333        assert_eq!(id.as_str(), "doc1");
334
335        let retrieved = collection.get(&id).unwrap();
336        assert_eq!(retrieved.get("name").and_then(|v| v.as_str()), Some("Alice"));
337    }
338
339    #[test]
340    fn test_duplicate_id() {
341        let collection = Collection::new("test");
342
343        let doc1 = Document::with_id("same-id");
344        let doc2 = Document::with_id("same-id");
345
346        collection.insert(doc1).unwrap();
347        let result = collection.insert(doc2);
348
349        assert!(matches!(result, Err(CollectionError::DuplicateId(_))));
350    }
351
352    #[test]
353    fn test_update() {
354        let collection = Collection::new("test");
355
356        let mut doc = Document::with_id("doc1");
357        doc.set("count", 1i64);
358        collection.insert(doc).unwrap();
359
360        let mut updated = Document::with_id("doc1");
361        updated.set("count", 2i64);
362        collection.update(&DocumentId::new("doc1"), updated).unwrap();
363
364        let retrieved = collection.get(&DocumentId::new("doc1")).unwrap();
365        assert_eq!(retrieved.get("count").and_then(|v| v.as_i64()), Some(2));
366    }
367
368    #[test]
369    fn test_delete() {
370        let collection = Collection::new("test");
371
372        let doc = Document::with_id("doc1");
373        collection.insert(doc).unwrap();
374
375        assert!(collection.contains(&DocumentId::new("doc1")));
376
377        collection.delete(&DocumentId::new("doc1")).unwrap();
378        assert!(!collection.contains(&DocumentId::new("doc1")));
379    }
380
381    #[test]
382    fn test_find() {
383        let collection = Collection::new("test");
384
385        for i in 0..10 {
386            let mut doc = Document::new();
387            doc.set("value", i as i64);
388            doc.set("even", i % 2 == 0);
389            collection.insert(doc).unwrap();
390        }
391
392        let query = QueryBuilder::new().eq("even", true).build();
393        let result = collection.find(&query);
394
395        assert_eq!(result.documents.len(), 5);
396    }
397}
aegis_document/collection.rs

aegis_document/
collection.rs
