Skip to main content

aegis_document/
engine.rs

1//! Aegis Document Engine
2//!
3//! Core engine that coordinates all document store operations.
4//!
5//! @version 0.1.0
6//! @author AutomataNexus Development Team
7
8use crate::collection::{Collection, CollectionError};
9use crate::index::IndexType;
10use crate::query::{Query, QueryResult};
11use crate::types::{Document, DocumentId};
12use crate::validation::Schema;
13use std::collections::HashMap;
14use std::sync::RwLock;
15
16// =============================================================================
17// Document Engine Configuration
18// =============================================================================
19
20/// Configuration for the document engine.
21#[derive(Debug, Clone)]
22pub struct EngineConfig {
23    pub max_document_size: usize,
24    pub max_collections: usize,
25    pub default_index_type: IndexType,
26    pub validate_on_insert: bool,
27}
28
29impl Default for EngineConfig {
30    fn default() -> Self {
31        Self {
32            max_document_size: 16 * 1024 * 1024, // 16MB
33            max_collections: 1000,
34            default_index_type: IndexType::Hash,
35            validate_on_insert: true,
36        }
37    }
38}
39
40// =============================================================================
41// Document Engine
42// =============================================================================
43
44/// The main document storage and query engine.
45pub struct DocumentEngine {
46    config: EngineConfig,
47    collections: RwLock<HashMap<String, Collection>>,
48    stats: RwLock<EngineStats>,
49}
50
51impl DocumentEngine {
52    /// Create a new document engine with default configuration.
53    pub fn new() -> Self {
54        Self::with_config(EngineConfig::default())
55    }
56
57    /// Create a new document engine with custom configuration.
58    pub fn with_config(config: EngineConfig) -> Self {
59        Self {
60            config,
61            collections: RwLock::new(HashMap::new()),
62            stats: RwLock::new(EngineStats::default()),
63        }
64    }
65
66    // -------------------------------------------------------------------------
67    // Collection Management
68    // -------------------------------------------------------------------------
69
70    /// Create a new collection.
71    pub fn create_collection(&self, name: impl Into<String>) -> Result<(), EngineError> {
72        let name = name.into();
73        let mut collections = self
74            .collections
75            .write()
76            .expect("collections RwLock poisoned");
77
78        if collections.len() >= self.config.max_collections {
79            return Err(EngineError::TooManyCollections);
80        }
81
82        if collections.contains_key(&name) {
83            return Err(EngineError::CollectionExists(name));
84        }
85
86        collections.insert(name.clone(), Collection::new(name));
87        Ok(())
88    }
89
90    /// Create a collection with a schema.
91    pub fn create_collection_with_schema(
92        &self,
93        name: impl Into<String>,
94        schema: Schema,
95    ) -> Result<(), EngineError> {
96        let name = name.into();
97        let mut collections = self
98            .collections
99            .write()
100            .expect("collections RwLock poisoned");
101
102        if collections.len() >= self.config.max_collections {
103            return Err(EngineError::TooManyCollections);
104        }
105
106        if collections.contains_key(&name) {
107            return Err(EngineError::CollectionExists(name));
108        }
109
110        collections.insert(name.clone(), Collection::with_schema(name, schema));
111        Ok(())
112    }
113
114    /// Drop a collection.
115    pub fn drop_collection(&self, name: &str) -> Result<(), EngineError> {
116        let mut collections = self
117            .collections
118            .write()
119            .expect("collections RwLock poisoned");
120
121        if collections.remove(name).is_none() {
122            return Err(EngineError::CollectionNotFound(name.to_string()));
123        }
124
125        Ok(())
126    }
127
128    /// List all collection names.
129    pub fn list_collections(&self) -> Vec<String> {
130        let collections = self
131            .collections
132            .read()
133            .expect("collections RwLock poisoned");
134        collections.keys().cloned().collect()
135    }
136
137    /// Check if a collection exists.
138    pub fn collection_exists(&self, name: &str) -> bool {
139        let collections = self
140            .collections
141            .read()
142            .expect("collections RwLock poisoned");
143        collections.contains_key(name)
144    }
145
146    /// Get collection statistics.
147    pub fn collection_stats(&self, name: &str) -> Option<CollectionStats> {
148        let collections = self
149            .collections
150            .read()
151            .expect("collections RwLock poisoned");
152        collections.get(name).map(|c| CollectionStats {
153            name: name.to_string(),
154            document_count: c.count(),
155            index_count: c.index_names().len(),
156        })
157    }
158
159    // -------------------------------------------------------------------------
160    // Document Operations
161    // -------------------------------------------------------------------------
162
163    /// Insert a document into a collection.
164    pub fn insert(&self, collection: &str, doc: Document) -> Result<DocumentId, EngineError> {
165        let collections = self
166            .collections
167            .read()
168            .expect("collections RwLock poisoned");
169        let coll = collections
170            .get(collection)
171            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
172
173        let id = coll.insert(doc).map_err(EngineError::Collection)?;
174
175        drop(collections);
176
177        {
178            let mut stats = self.stats.write().expect("stats RwLock poisoned");
179            stats.documents_inserted += 1;
180        }
181
182        Ok(id)
183    }
184
185    /// Insert multiple documents.
186    pub fn insert_many(
187        &self,
188        collection: &str,
189        docs: Vec<Document>,
190    ) -> Result<Vec<DocumentId>, EngineError> {
191        let collections = self
192            .collections
193            .read()
194            .expect("collections RwLock poisoned");
195        let coll = collections
196            .get(collection)
197            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
198
199        let count = docs.len();
200        let ids = coll.insert_many(docs).map_err(EngineError::Collection)?;
201
202        drop(collections);
203
204        {
205            let mut stats = self.stats.write().expect("stats RwLock poisoned");
206            stats.documents_inserted += count as u64;
207        }
208
209        Ok(ids)
210    }
211
212    /// Get a document by ID.
213    pub fn get(&self, collection: &str, id: &DocumentId) -> Result<Option<Document>, EngineError> {
214        let collections = self
215            .collections
216            .read()
217            .expect("collections RwLock poisoned");
218        let coll = collections
219            .get(collection)
220            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
221
222        Ok(coll.get(id))
223    }
224
225    /// Update a document.
226    pub fn update(
227        &self,
228        collection: &str,
229        id: &DocumentId,
230        doc: Document,
231    ) -> Result<(), EngineError> {
232        let collections = self
233            .collections
234            .read()
235            .expect("collections RwLock poisoned");
236        let coll = collections
237            .get(collection)
238            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
239
240        coll.update(id, doc).map_err(EngineError::Collection)?;
241
242        drop(collections);
243
244        {
245            let mut stats = self.stats.write().expect("stats RwLock poisoned");
246            stats.documents_updated += 1;
247        }
248
249        Ok(())
250    }
251
252    /// Delete a document.
253    pub fn delete(&self, collection: &str, id: &DocumentId) -> Result<Document, EngineError> {
254        let collections = self
255            .collections
256            .read()
257            .expect("collections RwLock poisoned");
258        let coll = collections
259            .get(collection)
260            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
261
262        let doc = coll.delete(id).map_err(EngineError::Collection)?;
263
264        drop(collections);
265
266        {
267            let mut stats = self.stats.write().expect("stats RwLock poisoned");
268            stats.documents_deleted += 1;
269        }
270
271        Ok(doc)
272    }
273
274    // -------------------------------------------------------------------------
275    // Query Operations
276    // -------------------------------------------------------------------------
277
278    /// Find documents matching a query.
279    pub fn find(&self, collection: &str, query: &Query) -> Result<QueryResult, EngineError> {
280        let collections = self
281            .collections
282            .read()
283            .expect("collections RwLock poisoned");
284        let coll = collections
285            .get(collection)
286            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
287
288        let result = coll.find(query);
289
290        drop(collections);
291
292        {
293            let mut stats = self.stats.write().expect("stats RwLock poisoned");
294            stats.queries_executed += 1;
295        }
296
297        Ok(result)
298    }
299
300    /// Find one document matching a query.
301    pub fn find_one(
302        &self,
303        collection: &str,
304        query: &Query,
305    ) -> Result<Option<Document>, EngineError> {
306        let collections = self
307            .collections
308            .read()
309            .expect("collections RwLock poisoned");
310        let coll = collections
311            .get(collection)
312            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
313
314        Ok(coll.find_one(query))
315    }
316
317    /// Count documents matching a query.
318    pub fn count(&self, collection: &str, query: &Query) -> Result<usize, EngineError> {
319        let collections = self
320            .collections
321            .read()
322            .expect("collections RwLock poisoned");
323        let coll = collections
324            .get(collection)
325            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
326
327        Ok(coll.count_matching(query))
328    }
329
330    // -------------------------------------------------------------------------
331    // Index Operations
332    // -------------------------------------------------------------------------
333
334    /// Create an index on a collection field.
335    pub fn create_index(
336        &self,
337        collection: &str,
338        field: impl Into<String>,
339        index_type: IndexType,
340    ) -> Result<(), EngineError> {
341        let collections = self
342            .collections
343            .read()
344            .expect("collections RwLock poisoned");
345        let coll = collections
346            .get(collection)
347            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
348
349        coll.create_index(field, index_type);
350        Ok(())
351    }
352
353    /// Drop an index.
354    pub fn drop_index(&self, collection: &str, field: &str) -> Result<(), EngineError> {
355        let collections = self
356            .collections
357            .read()
358            .expect("collections RwLock poisoned");
359        let coll = collections
360            .get(collection)
361            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
362
363        coll.drop_index(field);
364        Ok(())
365    }
366
367    /// List indexes on a collection.
368    pub fn list_indexes(&self, collection: &str) -> Result<Vec<String>, EngineError> {
369        let collections = self
370            .collections
371            .read()
372            .expect("collections RwLock poisoned");
373        let coll = collections
374            .get(collection)
375            .ok_or_else(|| EngineError::CollectionNotFound(collection.to_string()))?;
376
377        Ok(coll.index_names())
378    }
379
380    // -------------------------------------------------------------------------
381    // Statistics
382    // -------------------------------------------------------------------------
383
384    /// Get engine statistics.
385    pub fn stats(&self) -> EngineStats {
386        let stats = self.stats.read().expect("stats RwLock poisoned");
387        stats.clone()
388    }
389
390    /// Reset statistics.
391    pub fn reset_stats(&self) {
392        let mut stats = self.stats.write().expect("stats RwLock poisoned");
393        *stats = EngineStats::default();
394    }
395}
396
397impl Default for DocumentEngine {
398    fn default() -> Self {
399        Self::new()
400    }
401}
402
403// =============================================================================
404// Engine Statistics
405// =============================================================================
406
407/// Statistics for the document engine.
408#[derive(Debug, Clone, Default)]
409pub struct EngineStats {
410    pub documents_inserted: u64,
411    pub documents_updated: u64,
412    pub documents_deleted: u64,
413    pub queries_executed: u64,
414}
415
416/// Statistics for a collection.
417#[derive(Debug, Clone)]
418pub struct CollectionStats {
419    pub name: String,
420    pub document_count: usize,
421    pub index_count: usize,
422}
423
424// =============================================================================
425// Engine Error
426// =============================================================================
427
428/// Errors that can occur in the document engine.
429#[derive(Debug, Clone)]
430pub enum EngineError {
431    CollectionExists(String),
432    CollectionNotFound(String),
433    TooManyCollections,
434    Collection(CollectionError),
435    DocumentTooLarge,
436}
437
438impl std::fmt::Display for EngineError {
439    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
440        match self {
441            Self::CollectionExists(name) => write!(f, "Collection already exists: {}", name),
442            Self::CollectionNotFound(name) => write!(f, "Collection not found: {}", name),
443            Self::TooManyCollections => write!(f, "Maximum number of collections reached"),
444            Self::Collection(err) => write!(f, "Collection error: {}", err),
445            Self::DocumentTooLarge => write!(f, "Document exceeds maximum size"),
446        }
447    }
448}
449
450impl std::error::Error for EngineError {}
451
452// =============================================================================
453// Tests
454// =============================================================================
455
456#[cfg(test)]
457mod tests {
458    use super::*;
459    use crate::query::QueryBuilder;
460
461    #[test]
462    fn test_engine_creation() {
463        let engine = DocumentEngine::new();
464        assert!(engine.list_collections().is_empty());
465    }
466
467    #[test]
468    fn test_collection_management() {
469        let engine = DocumentEngine::new();
470
471        engine.create_collection("users").unwrap();
472        assert!(engine.collection_exists("users"));
473
474        let collections = engine.list_collections();
475        assert_eq!(collections.len(), 1);
476        assert!(collections.contains(&"users".to_string()));
477
478        engine.drop_collection("users").unwrap();
479        assert!(!engine.collection_exists("users"));
480    }
481
482    #[test]
483    fn test_document_crud() {
484        let engine = DocumentEngine::new();
485        engine.create_collection("test").unwrap();
486
487        let mut doc = Document::with_id("doc1");
488        doc.set("name", "Alice");
489        doc.set("age", 30i64);
490
491        let id = engine.insert("test", doc).unwrap();
492        assert_eq!(id.as_str(), "doc1");
493
494        let retrieved = engine.get("test", &id).unwrap().unwrap();
495        assert_eq!(
496            retrieved.get("name").and_then(|v| v.as_str()),
497            Some("Alice")
498        );
499
500        let mut updated = Document::with_id("doc1");
501        updated.set("name", "Alice Smith");
502        updated.set("age", 31i64);
503        engine.update("test", &id, updated).unwrap();
504
505        let retrieved = engine.get("test", &id).unwrap().unwrap();
506        assert_eq!(
507            retrieved.get("name").and_then(|v| v.as_str()),
508            Some("Alice Smith")
509        );
510
511        engine.delete("test", &id).unwrap();
512        assert!(engine.get("test", &id).unwrap().is_none());
513    }
514
515    #[test]
516    fn test_query() {
517        let engine = DocumentEngine::new();
518        engine.create_collection("products").unwrap();
519
520        for i in 0..10 {
521            let mut doc = Document::new();
522            doc.set("name", format!("Product {}", i));
523            doc.set("price", (i * 10) as i64);
524            doc.set("in_stock", i % 2 == 0);
525            engine.insert("products", doc).unwrap();
526        }
527
528        let query = QueryBuilder::new().eq("in_stock", true).build();
529        let result = engine.find("products", &query).unwrap();
530        assert_eq!(result.count(), 5);
531
532        let query = QueryBuilder::new().gt("price", 50i64).build();
533        let result = engine.find("products", &query).unwrap();
534        assert_eq!(result.count(), 4);
535    }
536
537    #[test]
538    fn test_index() {
539        let engine = DocumentEngine::new();
540        engine.create_collection("items").unwrap();
541
542        engine
543            .create_index("items", "category", IndexType::Hash)
544            .unwrap();
545
546        let indexes = engine.list_indexes("items").unwrap();
547        assert!(indexes.contains(&"category".to_string()));
548
549        engine.drop_index("items", "category").unwrap();
550        let indexes = engine.list_indexes("items").unwrap();
551        assert!(!indexes.contains(&"category".to_string()));
552    }
553
554    #[test]
555    fn test_stats() {
556        let engine = DocumentEngine::new();
557        engine.create_collection("test").unwrap();
558
559        for _ in 0..5 {
560            let doc = Document::new();
561            engine.insert("test", doc).unwrap();
562        }
563
564        let stats = engine.stats();
565        assert_eq!(stats.documents_inserted, 5);
566    }
567}