rust_zero_core/
stores_mongo.rs

1//! MongoDB collections, transactions, instrumentation, and cache-aside records.
2//!
3//! ```no_run
4//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
5//! use rust_zero_core::{MongoStore, MongoStoreConfig};
6//! let store = MongoStore::connect(MongoStoreConfig::new(
7//!     "mongodb://127.0.0.1:27017", "service",
8//! )).await?;
9//! store.health_check().await?;
10//! # Ok(())
11//! # }
12//! ```
13
14use mongodb::{
15    bson::{doc, Document},
16    error::Error as DriverError,
17    options::ClientOptions,
18    results::{DeleteResult, InsertManyResult, InsertOneResult, UpdateResult},
19    Client, ClientSession, Collection, Database,
20};
21use serde::Serialize;
22use std::{
23    error::Error,
24    fmt,
25    future::Future,
26    hash::Hash,
27    sync::{
28        atomic::{AtomicU64, Ordering},
29        Mutex,
30    },
31    time::{Duration, Instant},
32};
33
34use crate::{
35    cache::jittered_ttl, CacheStats, CounterVec, HistogramOptions, HistogramVec, MemoryCache,
36    Metrics, MetricsError, SingleFlight, SingleFlightError, VectorOptions,
37};
38
39#[cfg(feature = "telemetry")]
40use crate::{TelemetrySpan, TelemetrySpanKind};
41
42/// Cardinality-bounded MongoDB operation metrics used by [`MongoStore`]'s typed helpers.
43#[derive(Clone)]
44pub struct MongoStoreMetrics {
45    operations: CounterVec,
46    duration: HistogramVec,
47}
48
49impl MongoStoreMetrics {
50    pub fn register(metrics: &Metrics) -> Result<Self, MetricsError> {
51        let labels = ["operation", "kind", "outcome"];
52        Ok(Self {
53            operations: metrics.counter_vec(
54                VectorOptions::new("operations_total", "Completed MongoDB store operations")
55                    .with_namespace("rust_zero")
56                    .with_subsystem("mongo")
57                    .with_labels(labels),
58            )?,
59            duration: metrics.histogram_vec(
60                HistogramOptions::new(
61                    "operation_duration_seconds",
62                    "MongoDB store operation latency",
63                )
64                .with_vector_options(
65                    VectorOptions::new(
66                        "operation_duration_seconds",
67                        "MongoDB store operation latency",
68                    )
69                    .with_namespace("rust_zero")
70                    .with_subsystem("mongo")
71                    .with_labels(labels),
72                ),
73            )?,
74        })
75    }
76
77    fn observe(&self, operation: &str, kind: MongoOperationKind, outcome: &str, elapsed: Duration) {
78        let labels = [operation, kind.as_str(), outcome];
79        let _ = self.operations.inc(&labels);
80        let _ = self.duration.observe(elapsed.as_secs_f64(), &labels);
81    }
82}
83
84#[derive(Debug, Clone, Copy)]
85enum MongoOperationKind {
86    Query,
87    Execute,
88    BulkInsert,
89}
90
91impl MongoOperationKind {
92    fn as_str(self) -> &'static str {
93        match self {
94            Self::Query => "query",
95            Self::Execute => "execute",
96            Self::BulkInsert => "bulk_insert",
97        }
98    }
99}
100
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct MongoStoreConfig {
103    pub uri: String,
104    pub database: String,
105    pub application_name: Option<String>,
106    pub min_pool_size: Option<u32>,
107    pub max_pool_size: Option<u32>,
108    pub connect_timeout: Option<Duration>,
109    pub server_selection_timeout: Option<Duration>,
110}
111
112impl MongoStoreConfig {
113    pub fn new(uri: impl Into<String>, database: impl Into<String>) -> Self {
114        Self {
115            uri: uri.into(),
116            database: database.into(),
117            application_name: None,
118            min_pool_size: None,
119            max_pool_size: None,
120            connect_timeout: Some(Duration::from_secs(10)),
121            server_selection_timeout: Some(Duration::from_secs(10)),
122        }
123    }
124
125    pub fn with_application_name(mut self, name: impl Into<String>) -> Self {
126        self.application_name = Some(name.into());
127        self
128    }
129
130    pub fn with_pool_size(mut self, min: u32, max: u32) -> Self {
131        assert!(max > 0, "MongoDB maximum pool size must be positive");
132        assert!(
133            min <= max,
134            "MongoDB minimum pool size cannot exceed maximum"
135        );
136        self.min_pool_size = Some(min);
137        self.max_pool_size = Some(max);
138        self
139    }
140
141    pub fn with_timeouts(mut self, connect: Duration, server_selection: Duration) -> Self {
142        assert!(
143            !connect.is_zero(),
144            "MongoDB connect timeout must be positive"
145        );
146        assert!(
147            !server_selection.is_zero(),
148            "MongoDB server selection timeout must be positive"
149        );
150        self.connect_timeout = Some(connect);
151        self.server_selection_timeout = Some(server_selection);
152        self
153    }
154}
155
156/// Settings for the in-process cache used by [`CachedMongoCollection`].
157#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158pub struct MongoCacheConfig {
159    pub capacity: usize,
160    pub ttl: Duration,
161    pub not_found_ttl: Option<Duration>,
162    pub ttl_jitter: Duration,
163}
164
165impl MongoCacheConfig {
166    pub fn new(capacity: usize, ttl: Duration) -> Self {
167        assert!(capacity > 0, "MongoDB cache capacity must be positive");
168        assert!(!ttl.is_zero(), "MongoDB cache TTL must be positive");
169        Self {
170            capacity,
171            ttl,
172            not_found_ttl: Some(ttl),
173            ttl_jitter: Duration::ZERO,
174        }
175    }
176
177    /// Controls negative caching. `None` makes every missing record hit MongoDB.
178    pub fn with_not_found_ttl(mut self, ttl: Option<Duration>) -> Self {
179        assert!(
180            ttl.is_none_or(|ttl| !ttl.is_zero()),
181            "MongoDB not-found cache TTL must be positive"
182        );
183        self.not_found_ttl = ttl;
184        self
185    }
186
187    /// Randomizes each positive and negative expiry by up to `jitter`.
188    pub fn with_ttl_jitter(mut self, jitter: Duration) -> Self {
189        self.ttl_jitter = jitter;
190        self
191    }
192}
193
194#[derive(Debug)]
195pub enum MongoStoreError {
196    InvalidDatabase,
197    InvalidBatchSize,
198    NotFound { entity: String },
199    Driver(DriverError),
200}
201
202impl MongoStoreError {
203    pub fn not_found(entity: impl Into<String>) -> Self {
204        Self::NotFound {
205            entity: entity.into(),
206        }
207    }
208
209    pub fn is_not_found(&self) -> bool {
210        matches!(self, Self::NotFound { .. })
211    }
212}
213
214impl fmt::Display for MongoStoreError {
215    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
216        match self {
217            Self::InvalidDatabase => formatter.write_str("MongoDB database name cannot be empty"),
218            Self::InvalidBatchSize => {
219                formatter.write_str("MongoDB bulk batch size must be positive")
220            }
221            Self::NotFound { entity } => write!(formatter, "{entity} not found"),
222            Self::Driver(error) => write!(formatter, "MongoDB operation failed: {error}"),
223        }
224    }
225}
226
227impl Error for MongoStoreError {
228    fn source(&self) -> Option<&(dyn Error + 'static)> {
229        match self {
230            Self::InvalidDatabase | Self::InvalidBatchSize | Self::NotFound { .. } => None,
231            Self::Driver(error) => Some(error),
232        }
233    }
234}
235
236impl From<DriverError> for MongoStoreError {
237    fn from(error: DriverError) -> Self {
238        Self::Driver(error)
239    }
240}
241
242/// A reusable MongoDB client and selected database.
243#[derive(Clone)]
244pub struct MongoStore {
245    client: Client,
246    database: Database,
247    metrics: Option<MongoStoreMetrics>,
248}
249
250impl fmt::Debug for MongoStore {
251    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
252        formatter
253            .debug_struct("MongoStore")
254            .field("client", &self.client)
255            .field("database", &self.database)
256            .field("instrumented", &self.metrics.is_some())
257            .finish()
258    }
259}
260
261impl MongoStore {
262    pub async fn connect(config: MongoStoreConfig) -> Result<Self, MongoStoreError> {
263        if config.database.trim().is_empty() {
264            return Err(MongoStoreError::InvalidDatabase);
265        }
266
267        let mut options = ClientOptions::parse(&config.uri).await?;
268        options.app_name = config.application_name;
269        options.min_pool_size = config.min_pool_size;
270        options.max_pool_size = config.max_pool_size;
271        options.connect_timeout = config.connect_timeout;
272        options.server_selection_timeout = config.server_selection_timeout;
273        let client = Client::with_options(options)?;
274        let database = client.database(&config.database);
275        Ok(Self {
276            client,
277            database,
278            metrics: None,
279        })
280    }
281
282    pub fn from_client(client: Client, database: impl AsRef<str>) -> Result<Self, MongoStoreError> {
283        let database = database.as_ref();
284        if database.trim().is_empty() {
285            return Err(MongoStoreError::InvalidDatabase);
286        }
287        Ok(Self {
288            database: client.database(database),
289            client,
290            metrics: None,
291        })
292    }
293
294    /// Installs metrics for operations run through the typed helper methods.
295    pub fn with_metrics(mut self, metrics: MongoStoreMetrics) -> Self {
296        self.metrics = Some(metrics);
297        self
298    }
299
300    pub fn client(&self) -> &Client {
301        &self.client
302    }
303
304    pub fn database(&self) -> &Database {
305        &self.database
306    }
307
308    pub fn collection<T>(&self, name: impl AsRef<str>) -> Collection<T>
309    where
310        T: Send + Sync,
311    {
312        self.database.collection(name.as_ref())
313    }
314
315    pub fn cached_collection<K, T>(
316        &self,
317        name: impl AsRef<str>,
318        config: MongoCacheConfig,
319    ) -> CachedMongoCollection<K, T>
320    where
321        K: Clone + Eq + Hash,
322        T: Clone + Send + Sync,
323    {
324        CachedMongoCollection::new(self.collection(name), config)
325    }
326
327    /// Creates a cached collection with distinct primary- and secondary-key types.
328    pub fn cached_indexed_collection<K, I, T>(
329        &self,
330        name: impl AsRef<str>,
331        config: MongoCacheConfig,
332    ) -> CachedMongoCollection<K, T, I>
333    where
334        K: Clone + Eq + Hash,
335        I: Clone + Eq + Hash,
336        T: Clone + Send + Sync,
337    {
338        CachedMongoCollection::new(self.collection(name), config)
339    }
340
341    pub async fn health_check(&self) -> Result<(), MongoStoreError> {
342        self.database.run_command(doc! { "ping": 1 }).await?;
343        Ok(())
344    }
345
346    /// Runs a typed collection query with consistent metrics, tracing, and error conversion.
347    pub async fn query<T, R, F, Fut>(
348        &self,
349        operation: &'static str,
350        collection: impl AsRef<str>,
351        query: F,
352    ) -> Result<R, MongoStoreError>
353    where
354        T: Send + Sync,
355        F: FnOnce(Collection<T>) -> Fut,
356        Fut: Future<Output = Result<R, DriverError>>,
357    {
358        let collection = self.collection(collection);
359        self.instrument(operation, MongoOperationKind::Query, async move {
360            query(collection).await.map_err(MongoStoreError::Driver)
361        })
362        .await
363    }
364
365    /// Runs an optional query and converts an absent document into a stable not-found error.
366    pub async fn query_one<T, F, Fut>(
367        &self,
368        operation: &'static str,
369        collection: impl AsRef<str>,
370        entity: impl Into<String>,
371        query: F,
372    ) -> Result<T, MongoStoreError>
373    where
374        T: Send + Sync,
375        F: FnOnce(Collection<T>) -> Fut,
376        Fut: Future<Output = Result<Option<T>, DriverError>>,
377    {
378        let collection = self.collection(collection);
379        let entity = entity.into();
380        self.instrument(operation, MongoOperationKind::Query, async move {
381            query(collection)
382                .await
383                .map_err(MongoStoreError::Driver)?
384                .ok_or_else(|| MongoStoreError::not_found(entity))
385        })
386        .await
387    }
388
389    /// Runs a typed collection mutation with consistent metrics and tracing.
390    pub async fn execute<T, R, F, Fut>(
391        &self,
392        operation: &'static str,
393        collection: impl AsRef<str>,
394        execute: F,
395    ) -> Result<R, MongoStoreError>
396    where
397        T: Send + Sync,
398        F: FnOnce(Collection<T>) -> Fut,
399        Fut: Future<Output = Result<R, DriverError>>,
400    {
401        let collection = self.collection(collection);
402        self.instrument(operation, MongoOperationKind::Execute, async move {
403            execute(collection).await.map_err(MongoStoreError::Driver)
404        })
405        .await
406    }
407
408    /// Inserts documents with MongoDB's native `insert_many` in bounded batches.
409    pub async fn bulk_insert<T>(
410        &self,
411        operation: &'static str,
412        collection: impl AsRef<str>,
413        items: impl IntoIterator<Item = T>,
414        batch_size: usize,
415    ) -> Result<Vec<InsertManyResult>, MongoStoreError>
416    where
417        T: Serialize + Send + Sync,
418    {
419        if batch_size == 0 {
420            return Err(MongoStoreError::InvalidBatchSize);
421        }
422
423        let collection = self.collection::<T>(collection);
424        let mut items = items.into_iter();
425        self.instrument(operation, MongoOperationKind::BulkInsert, async move {
426            let mut results = Vec::new();
427            loop {
428                let batch: Vec<_> = items.by_ref().take(batch_size).collect();
429                if batch.is_empty() {
430                    return Ok(results);
431                }
432                results.push(collection.insert_many(batch).await?);
433            }
434        })
435        .await
436    }
437
438    async fn instrument<R, Fut>(
439        &self,
440        operation: &'static str,
441        kind: MongoOperationKind,
442        future: Fut,
443    ) -> Result<R, MongoStoreError>
444    where
445        Fut: Future<Output = Result<R, MongoStoreError>>,
446    {
447        let started = Instant::now();
448        #[cfg(feature = "telemetry")]
449        let span = TelemetrySpan::start(
450            format!("mongo.{operation}"),
451            TelemetrySpanKind::Client,
452            None,
453            [
454                ("db.operation.name", operation.to_owned()),
455                ("rust_zero.mongo.kind", kind.as_str().to_owned()),
456            ],
457        );
458        let result = future.await;
459        let outcome = mongo_outcome(&result);
460        if let Some(metrics) = &self.metrics {
461            metrics.observe(operation, kind, outcome, started.elapsed());
462        }
463        #[cfg(feature = "telemetry")]
464        if let Err(error) = &result {
465            span.set_error(error.to_string());
466        }
467        result
468    }
469
470    /// Starts a client session and transaction. Pass the returned session to each operation.
471    pub async fn begin(&self) -> Result<ClientSession, MongoStoreError> {
472        let mut session = self.client.start_session().await?;
473        session.start_transaction().await?;
474        Ok(session)
475    }
476}
477
478fn mongo_outcome<T>(result: &Result<T, MongoStoreError>) -> &'static str {
479    match result {
480        Ok(_) => "success",
481        Err(MongoStoreError::NotFound { .. }) => "not_found",
482        Err(
483            MongoStoreError::InvalidDatabase
484            | MongoStoreError::InvalidBatchSize
485            | MongoStoreError::Driver(_),
486        ) => "error",
487    }
488}
489
490/// A typed MongoDB collection with bounded positive/negative record caching.
491pub struct CachedMongoCollection<K, T, I = K>
492where
493    T: Send + Sync,
494{
495    collection: Collection<T>,
496    cache: MemoryCache<K, Option<T>>,
497    indexes: MemoryCache<I, Option<K>>,
498    flights: SingleFlight<K, Option<T>, DriverError>,
499    index_flights: SingleFlight<I, Option<(K, T)>, DriverError>,
500    ttl: Duration,
501    not_found_ttl: Option<Duration>,
502    ttl_jitter: Duration,
503    expiry_sequence: AtomicU64,
504    generation: AtomicU64,
505    cache_gate: Mutex<()>,
506}
507
508impl<K, T, I> CachedMongoCollection<K, T, I>
509where
510    K: Clone + Eq + Hash,
511    I: Clone + Eq + Hash,
512    T: Clone + Send + Sync,
513{
514    pub fn new(collection: Collection<T>, config: MongoCacheConfig) -> Self {
515        Self {
516            collection,
517            cache: MemoryCache::new(config.capacity),
518            indexes: MemoryCache::new(config.capacity),
519            flights: SingleFlight::new(),
520            index_flights: SingleFlight::new(),
521            ttl: config.ttl,
522            not_found_ttl: config.not_found_ttl,
523            ttl_jitter: config.ttl_jitter,
524            expiry_sequence: AtomicU64::new(0),
525            generation: AtomicU64::new(0),
526            cache_gate: Mutex::new(()),
527        }
528    }
529
530    pub fn collection(&self) -> &Collection<T> {
531        &self.collection
532    }
533
534    pub async fn find<F, Fut>(
535        &self,
536        key: K,
537        query: F,
538    ) -> Result<Option<T>, SingleFlightError<DriverError>>
539    where
540        F: FnOnce(Collection<T>) -> Fut,
541        Fut: Future<Output = Result<Option<T>, DriverError>>,
542    {
543        if let Some(value) = self.cache.get(&key) {
544            return Ok(value);
545        }
546
547        self.flights
548            .execute(key.clone(), || async {
549                if let Some(value) = self.cache.get(&key) {
550                    return Ok(value);
551                }
552                let generation = self.generation.load(Ordering::Acquire);
553                let value = query(self.collection.clone()).await?;
554                let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
555                if self.generation.load(Ordering::Acquire) == generation {
556                    let base_ttl = if value.is_some() {
557                        Some(self.ttl)
558                    } else {
559                        self.not_found_ttl
560                    };
561                    if let Some(base_ttl) = base_ttl {
562                        let sequence = self.expiry_sequence.fetch_add(1, Ordering::Relaxed);
563                        self.cache.insert(
564                            key,
565                            value.clone(),
566                            jittered_ttl(base_ttl, self.ttl_jitter, sequence),
567                        );
568                    }
569                }
570                Ok(value)
571            })
572            .await
573    }
574
575    /// Returns a document selected by a secondary key while caching its primary-key mapping.
576    pub async fn find_by_index<F, Fut>(
577        &self,
578        index: I,
579        query: F,
580    ) -> Result<Option<T>, SingleFlightError<DriverError>>
581    where
582        F: FnOnce(Collection<T>) -> Fut,
583        Fut: Future<Output = Result<Option<(K, T)>, DriverError>>,
584    {
585        if let Some(primary) = self.indexes.get(&index) {
586            match primary {
587                Some(primary) => {
588                    if let Some(value) = self.cache.get(&primary) {
589                        return Ok(value);
590                    }
591                }
592                None => return Ok(None),
593            }
594        }
595
596        let loaded = self
597            .index_flights
598            .execute(index.clone(), || async {
599                if let Some(primary) = self.indexes.get(&index) {
600                    match primary {
601                        Some(primary) => {
602                            if let Some(Some(value)) = self.cache.get(&primary) {
603                                return Ok(Some((primary, value)));
604                            }
605                        }
606                        None => return Ok(None),
607                    }
608                }
609
610                let generation = self.generation.load(Ordering::Acquire);
611                let value = query(self.collection.clone()).await?;
612                let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
613                if self.generation.load(Ordering::Acquire) == generation {
614                    let base_ttl = if value.is_some() {
615                        Some(self.ttl)
616                    } else {
617                        self.not_found_ttl
618                    };
619                    if let Some(base_ttl) = base_ttl {
620                        let sequence = self.expiry_sequence.fetch_add(1, Ordering::Relaxed);
621                        let ttl = jittered_ttl(base_ttl, self.ttl_jitter, sequence);
622                        let primary = value.as_ref().map(|(primary, _)| primary.clone());
623                        self.indexes.insert(index, primary, ttl);
624                        if let Some((primary, value)) = &value {
625                            self.cache.insert(primary.clone(), Some(value.clone()), ttl);
626                        }
627                    }
628                }
629                Ok(value)
630            })
631            .await?;
632        Ok(loaded.map(|(_, value)| value))
633    }
634
635    pub async fn execute<PI, F, Fut, R>(&self, keys: PI, operation: F) -> Result<R, DriverError>
636    where
637        PI: IntoIterator<Item = K>,
638        F: FnOnce(Collection<T>) -> Fut,
639        Fut: Future<Output = Result<R, DriverError>>,
640    {
641        let result = operation(self.collection.clone()).await?;
642        self.invalidate_many(keys);
643        Ok(result)
644    }
645
646    /// Runs a mutation and invalidates primary keys plus explicitly affected secondary keys.
647    ///
648    /// Learned mappings for the primary keys are removed automatically. Pass changed/new index
649    /// keys as well so negative index entries cannot hide inserts or updates.
650    pub async fn execute_indexed<PI, SI, F, Fut, R>(
651        &self,
652        primary_keys: PI,
653        index_keys: SI,
654        operation: F,
655    ) -> Result<R, DriverError>
656    where
657        PI: IntoIterator<Item = K>,
658        SI: IntoIterator<Item = I>,
659        F: FnOnce(Collection<T>) -> Fut,
660        Fut: Future<Output = Result<R, DriverError>>,
661    {
662        let result = operation(self.collection.clone()).await?;
663        self.invalidate_related(primary_keys, index_keys);
664        Ok(result)
665    }
666
667    /// Inserts one document and invalidates any cached negative primary or secondary lookups.
668    pub async fn insert_one<SI>(
669        &self,
670        primary_key: K,
671        index_keys: SI,
672        document: T,
673    ) -> Result<InsertOneResult, DriverError>
674    where
675        SI: IntoIterator<Item = I>,
676        T: Serialize,
677    {
678        self.execute_indexed([primary_key], index_keys, move |collection| async move {
679            collection.insert_one(document).await
680        })
681        .await
682    }
683
684    /// Inserts documents in one native MongoDB operation and invalidates their cached keys.
685    pub async fn insert_many<PI, SI>(
686        &self,
687        primary_keys: PI,
688        index_keys: SI,
689        documents: Vec<T>,
690    ) -> Result<InsertManyResult, DriverError>
691    where
692        PI: IntoIterator<Item = K>,
693        SI: IntoIterator<Item = I>,
694        T: Serialize,
695    {
696        self.execute_indexed(primary_keys, index_keys, move |collection| async move {
697            collection.insert_many(documents).await
698        })
699        .await
700    }
701
702    /// Updates one document and invalidates all affected primary and secondary cache entries.
703    pub async fn update_one<PI, SI>(
704        &self,
705        primary_keys: PI,
706        index_keys: SI,
707        filter: Document,
708        update: Document,
709    ) -> Result<UpdateResult, DriverError>
710    where
711        PI: IntoIterator<Item = K>,
712        SI: IntoIterator<Item = I>,
713    {
714        self.execute_indexed(primary_keys, index_keys, move |collection| async move {
715            collection.update_one(filter, update).await
716        })
717        .await
718    }
719
720    /// Replaces one document and invalidates all affected primary and secondary cache entries.
721    pub async fn replace_one<PI, SI>(
722        &self,
723        primary_keys: PI,
724        index_keys: SI,
725        filter: Document,
726        replacement: T,
727    ) -> Result<UpdateResult, DriverError>
728    where
729        PI: IntoIterator<Item = K>,
730        SI: IntoIterator<Item = I>,
731        T: Serialize,
732    {
733        self.execute_indexed(primary_keys, index_keys, move |collection| async move {
734            collection.replace_one(filter, replacement).await
735        })
736        .await
737    }
738
739    /// Deletes one document and invalidates all affected primary and secondary cache entries.
740    pub async fn delete_one<PI, SI>(
741        &self,
742        primary_keys: PI,
743        index_keys: SI,
744        filter: Document,
745    ) -> Result<DeleteResult, DriverError>
746    where
747        PI: IntoIterator<Item = K>,
748        SI: IntoIterator<Item = I>,
749    {
750        self.execute_indexed(primary_keys, index_keys, move |collection| async move {
751            collection.delete_one(filter).await
752        })
753        .await
754    }
755
756    pub fn invalidate(&self, key: &K) -> bool {
757        let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
758        self.generation.fetch_add(1, Ordering::AcqRel);
759        let removed = self.cache.remove(key).is_some();
760        self.indexes
761            .remove_where(|_, primary| primary.as_ref() == Some(key));
762        removed
763    }
764
765    pub fn invalidate_many<PI>(&self, keys: PI)
766    where
767        PI: IntoIterator<Item = K>,
768    {
769        let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
770        self.generation.fetch_add(1, Ordering::AcqRel);
771        let keys: std::collections::HashSet<_> = keys.into_iter().collect();
772        for key in &keys {
773            self.cache.remove(key);
774        }
775        self.indexes
776            .remove_where(|_, primary| primary.as_ref().is_some_and(|key| keys.contains(key)));
777    }
778
779    pub fn invalidate_index(&self, index: &I) -> bool {
780        let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
781        self.generation.fetch_add(1, Ordering::AcqRel);
782        self.indexes.remove(index).is_some()
783    }
784
785    pub fn invalidate_related<PI, SI>(&self, primary_keys: PI, index_keys: SI)
786    where
787        PI: IntoIterator<Item = K>,
788        SI: IntoIterator<Item = I>,
789    {
790        let _guard = self.cache_gate.lock().expect("MongoDB cache gate poisoned");
791        self.generation.fetch_add(1, Ordering::AcqRel);
792        let primary_keys: std::collections::HashSet<_> = primary_keys.into_iter().collect();
793        for key in &primary_keys {
794            self.cache.remove(key);
795        }
796        self.indexes.remove_where(|_, primary| {
797            primary
798                .as_ref()
799                .is_some_and(|key| primary_keys.contains(key))
800        });
801        for index in index_keys {
802            self.indexes.remove(&index);
803        }
804    }
805
806    pub fn cache_stats(&self) -> CacheStats {
807        self.cache.stats()
808    }
809
810    pub fn index_cache_stats(&self) -> CacheStats {
811        self.indexes.stats()
812    }
813}
814
815#[cfg(test)]
816mod tests {
817    use super::*;
818    use serde::{Deserialize, Serialize};
819    use std::sync::atomic::{AtomicUsize, Ordering};
820
821    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
822    struct User {
823        #[serde(rename = "_id")]
824        id: i64,
825        name: String,
826    }
827
828    #[test]
829    fn cache_policy_supports_jitter_and_optional_negative_entries() {
830        let config = MongoCacheConfig::new(128, Duration::from_secs(30))
831            .with_not_found_ttl(None)
832            .with_ttl_jitter(Duration::from_secs(5));
833
834        assert_eq!(config.not_found_ttl, None);
835        assert_eq!(config.ttl_jitter, Duration::from_secs(5));
836    }
837
838    #[tokio::test]
839    async fn validates_configuration_and_caches_positive_and_negative_results() {
840        let error = MongoStore::connect(MongoStoreConfig::new("mongodb://127.0.0.1:27017", " "))
841            .await
842            .unwrap_err();
843        assert!(matches!(error, MongoStoreError::InvalidDatabase));
844
845        let store = MongoStore::connect(MongoStoreConfig::new(
846            "mongodb://127.0.0.1:27017",
847            "rust_zero_test",
848        ))
849        .await
850        .unwrap();
851        let cached = store.cached_collection::<i64, User>(
852            "users",
853            MongoCacheConfig::new(10, Duration::from_secs(30)),
854        );
855        let queries = AtomicUsize::new(0);
856
857        for _ in 0..2 {
858            let queries = &queries;
859            let user = cached
860                .find(7, move |_collection| async move {
861                    queries.fetch_add(1, Ordering::SeqCst);
862                    Ok(Some(User {
863                        id: 7,
864                        name: "Ada".to_owned(),
865                    }))
866                })
867                .await
868                .unwrap();
869            assert_eq!(user.unwrap().name, "Ada");
870        }
871        for _ in 0..2 {
872            let queries = &queries;
873            assert!(cached
874                .find(404, move |_collection| async move {
875                    queries.fetch_add(1, Ordering::SeqCst);
876                    Ok(None)
877                })
878                .await
879                .unwrap()
880                .is_none());
881        }
882        assert_eq!(queries.load(Ordering::SeqCst), 2);
883        assert!(cached.invalidate(&7));
884    }
885
886    #[tokio::test]
887    async fn typed_helpers_standardize_not_found_validate_batches_and_emit_metrics() {
888        let registry = Metrics::new();
889        let metrics = MongoStoreMetrics::register(&registry).unwrap();
890        let store = MongoStore::connect(MongoStoreConfig::new(
891            "mongodb://127.0.0.1:27017",
892            "rust_zero_test",
893        ))
894        .await
895        .unwrap()
896        .with_metrics(metrics);
897
898        let user = store
899            .query::<User, _, _, _>("find_user", "users", |_collection| async {
900                Ok(User {
901                    id: 7,
902                    name: "Ada".to_owned(),
903                })
904            })
905            .await
906            .unwrap();
907        assert_eq!(user.name, "Ada");
908
909        let missing = store
910            .query_one::<User, _, _>("find_user", "users", "user", |_collection| async {
911                Ok(None)
912            })
913            .await
914            .unwrap_err();
915        assert!(missing.is_not_found());
916        assert_eq!(missing.to_string(), "user not found");
917
918        assert!(matches!(
919            store.bulk_insert::<User>("invalid", "users", [], 0).await,
920            Err(MongoStoreError::InvalidBatchSize)
921        ));
922        let rendered = registry.render();
923        assert!(rendered.contains(
924            "rust_zero_mongo_operations_total{operation=\"find_user\",kind=\"query\",outcome=\"success\"} 1"
925        ));
926        assert!(rendered.contains(
927            "rust_zero_mongo_operations_total{operation=\"find_user\",kind=\"query\",outcome=\"not_found\"} 1"
928        ));
929    }
930
931    #[tokio::test]
932    async fn secondary_indexes_share_documents_and_follow_primary_invalidation() {
933        let store = MongoStore::connect(MongoStoreConfig::new(
934            "mongodb://127.0.0.1:27017",
935            "rust_zero_test",
936        ))
937        .await
938        .unwrap();
939        let cached = store.cached_indexed_collection::<i64, String, User>(
940            "users",
941            MongoCacheConfig::new(10, Duration::from_secs(30)),
942        );
943        let queries = AtomicUsize::new(0);
944
945        for _ in 0..2 {
946            let queries = &queries;
947            let user = cached
948                .find_by_index("ada@test".to_owned(), move |_collection| async move {
949                    queries.fetch_add(1, Ordering::SeqCst);
950                    Ok(Some((
951                        7,
952                        User {
953                            id: 7,
954                            name: "Ada".to_owned(),
955                        },
956                    )))
957                })
958                .await
959                .unwrap();
960            assert_eq!(user.unwrap().name, "Ada");
961        }
962        assert_eq!(queries.load(Ordering::SeqCst), 1);
963
964        assert!(cached.invalidate(&7));
965        let queries = &queries;
966        let user = cached
967            .find_by_index("ada@test".to_owned(), move |_collection| async move {
968                queries.fetch_add(1, Ordering::SeqCst);
969                Ok(Some((
970                    7,
971                    User {
972                        id: 7,
973                        name: "Grace".to_owned(),
974                    },
975                )))
976            })
977            .await
978            .unwrap();
979        assert_eq!(user.unwrap().name, "Grace");
980        assert_eq!(queries.load(Ordering::SeqCst), 2);
981        assert!(cached.index_cache_stats().insertions >= 2);
982    }
983
984    #[tokio::test]
985    async fn mongodb_integration_covers_health_crud_cache_and_transactions() {
986        let Ok(uri) = std::env::var("RUST_ZERO_MONGODB_URI") else {
987            return;
988        };
989        let database = format!("rust_zero_{}", std::process::id());
990        let store = MongoStore::connect(MongoStoreConfig::new(uri, &database))
991            .await
992            .unwrap();
993        store.health_check().await.unwrap();
994        let users = store.cached_collection::<i64, User>(
995            "users",
996            MongoCacheConfig::new(100, Duration::from_secs(30)),
997        );
998        users
999            .insert_one(
1000                7,
1001                std::iter::empty(),
1002                User {
1003                    id: 7,
1004                    name: "Ada".to_owned(),
1005                },
1006            )
1007            .await
1008            .unwrap();
1009        let user = users
1010            .find(7, |collection| async move {
1011                collection.find_one(doc! { "_id": 7_i64 }).await
1012            })
1013            .await
1014            .unwrap()
1015            .unwrap();
1016        assert_eq!(user.name, "Ada");
1017        users
1018            .update_one(
1019                [7],
1020                std::iter::empty(),
1021                doc! { "_id": 7_i64 },
1022                doc! { "$set": { "name": "Grace" } },
1023            )
1024            .await
1025            .unwrap();
1026        let user = users
1027            .find(7, |collection| async move {
1028                collection.find_one(doc! { "_id": 7_i64 }).await
1029            })
1030            .await
1031            .unwrap()
1032            .unwrap();
1033        assert_eq!(user.name, "Grace");
1034
1035        let bulk = store
1036            .bulk_insert(
1037                "insert_users",
1038                "users",
1039                (8_i64..=10).map(|id| User {
1040                    id,
1041                    name: format!("user-{id}"),
1042                }),
1043                2,
1044            )
1045            .await
1046            .unwrap();
1047        assert_eq!(bulk.len(), 2);
1048        assert_eq!(
1049            bulk.iter()
1050                .map(|result| result.inserted_ids.len())
1051                .sum::<usize>(),
1052            3
1053        );
1054
1055        if std::env::var_os("RUST_ZERO_MONGODB_TRANSACTIONS").is_some() {
1056            let mut transaction = store.begin().await.unwrap();
1057            users
1058                .collection()
1059                .delete_one(doc! { "_id": 7_i64 })
1060                .session(&mut transaction)
1061                .await
1062                .unwrap();
1063            transaction.abort_transaction().await.unwrap();
1064        }
1065        store.database().drop().await.unwrap();
1066    }
1067}
rust_zero_core/stores_mongo.rs

rust_zero_core/
stores_mongo.rs
