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