Skip to main content

faucet_source_mongodb_cdc/
stream.rs

1//! MongoDB change-stream executor.
2
3use crate::config::{MongoCdcSourceConfig, StartFrom};
4use crate::envelope::to_envelope;
5use crate::state::{Bookmark, state_key};
6use async_trait::async_trait;
7use faucet_core::{FaucetError, Source, Stream, StreamPage};
8use mongodb::Client;
9use mongodb::bson::{self, Document, Timestamp};
10use mongodb::change_stream::event::ResumeToken;
11use mongodb::options::{FullDocumentBeforeChangeType, FullDocumentType};
12use serde_json::Value;
13use std::collections::HashMap;
14use std::pin::Pin;
15use tokio::sync::Mutex;
16
17/// Effective start position resolved for a given fetch cycle.
18#[derive(Clone, Debug, PartialEq)]
19pub(crate) enum StartPosition {
20    /// Server default (current cluster time).
21    Now,
22    /// `start_at_operation_time`.
23    AtOperationTime(Timestamp),
24    /// `resume_after(token)`.
25    ResumeAfter(ResumeToken),
26    /// `start_after(token)` — required (not `resume_after`) when resuming from an
27    /// **invalidate** token, which MongoDB rejects for `resumeAfter` (#321 M3).
28    StartAfter(ResumeToken),
29}
30
31/// Resolve the effective start position from config + any persisted bookmark.
32///
33/// Precedence: explicit `resume_token`/`timestamp` win over a persisted
34/// bookmark; `now`/`earliest` yield to a persisted bookmark.
35pub(crate) fn resolve_start(
36    start_from: &StartFrom,
37    pending: Option<&Bookmark>,
38) -> Result<StartPosition, FaucetError> {
39    match start_from {
40        StartFrom::ResumeToken { token } => {
41            let b = Bookmark {
42                resume_token: token.clone(),
43                ..Default::default()
44            };
45            Ok(StartPosition::ResumeAfter(b.to_token()?))
46        }
47        StartFrom::Timestamp { timestamp_secs } => Ok(StartPosition::AtOperationTime(Timestamp {
48            time: *timestamp_secs,
49            increment: 0,
50        })),
51        StartFrom::Now | StartFrom::Earliest => {
52            if let Some(b) = pending {
53                // A bookmark captured from an invalidate event must resume with
54                // `start_after` — `resume_after` on an invalidate token is
55                // rejected by MongoDB and would wedge the pipeline (#321 M3).
56                if b.invalidate {
57                    Ok(StartPosition::StartAfter(b.to_token()?))
58                } else {
59                    Ok(StartPosition::ResumeAfter(b.to_token()?))
60                }
61            } else if matches!(start_from, StartFrom::Earliest) {
62                // Earliest retained oplog entry (errors at open if rolled).
63                Ok(StartPosition::AtOperationTime(Timestamp {
64                    time: 1,
65                    increment: 1,
66                }))
67            } else {
68                Ok(StartPosition::Now)
69            }
70        }
71    }
72}
73
74/// Enforce the in-memory staging cap before buffering one more change record.
75///
76/// `current_len` is the number of records already buffered this fetch cycle;
77/// `max_staged` is the configured cap (`None` = unbounded). Returns a typed
78/// [`FaucetError::Source`] when adding one more record would exceed the cap,
79/// so the run aborts cleanly instead of risking an OOM-kill on an unbounded
80/// `batch_size: 0` (or `fetch_all`) drain. Mirrors `postgres-cdc` /
81/// `mysql-cdc`'s `push_staged` guard.
82pub(crate) fn check_staging_cap(
83    current_len: usize,
84    max_staged: Option<usize>,
85) -> Result<(), FaucetError> {
86    if let Some(max) = max_staged
87        && current_len >= max
88    {
89        return Err(FaucetError::Source(format!(
90            "mongodb-cdc: in-memory change buffer exceeded max_staged_records ({max}); \
91             aborting to avoid unbounded memory growth. Raise max_staged_records or \
92             reduce batch_size so pages flush more often."
93        )));
94    }
95    Ok(())
96}
97
98/// A configured MongoDB CDC source.
99pub struct MongoCdcSource {
100    config: MongoCdcSourceConfig,
101    client: Client,
102    state_key_value: String,
103    pending_bookmark: Mutex<Option<Bookmark>>,
104}
105
106impl MongoCdcSource {
107    /// Connect, validate config + topology, and build the source.
108    pub async fn new(config: MongoCdcSourceConfig) -> Result<Self, FaucetError> {
109        config.validate()?;
110        let client = Client::with_uri_str(&config.connection_uri)
111            .await
112            .map_err(|e| FaucetError::Source(format!("MongoDB connection failed: {e}")))?;
113        ensure_changestream_capable(&client).await?;
114        if config.full_document_before_change != crate::config::FullDocumentBeforeChange::Off {
115            tracing::warn!(
116                "mongodb-cdc: full_document_before_change requires MongoDB 6.0+ and the \
117                 collection's changeStreamPreAndPostImages enabled; the server will error \
118                 at stream open if unavailable"
119            );
120        }
121        let state_key_value = state_key(&config.scope);
122        Ok(Self {
123            config,
124            client,
125            state_key_value,
126            pending_bookmark: Mutex::new(None),
127        })
128    }
129}
130
131/// Run `hello` and assert the deployment supports change streams.
132async fn ensure_changestream_capable(client: &Client) -> Result<(), FaucetError> {
133    let hello = client
134        .database("admin")
135        .run_command(bson::doc! { "hello": 1 })
136        .await
137        .map_err(|e| FaucetError::Source(format!("mongodb-cdc hello failed: {e}")))?;
138    if is_changestream_topology(&hello) {
139        Ok(())
140    } else {
141        Err(FaucetError::Source(
142            "mongodb-cdc: Change Streams require a replica set or sharded cluster; \
143             the target appears to be a standalone mongod"
144                .into(),
145        ))
146    }
147}
148
149/// Classify a `hello` response: replica-set (`setName`) or sharded (`msg=isdbgrid`).
150pub(crate) fn is_changestream_topology(hello: &Document) -> bool {
151    hello.get_str("setName").is_ok()
152        || hello
153            .get_str("msg")
154            .map(|m| m == "isdbgrid")
155            .unwrap_or(false)
156}
157
158/// Map the config post-image mode to the driver type (`None` = don't request).
159pub(crate) fn full_document_type(fd: crate::config::FullDocument) -> Option<FullDocumentType> {
160    use crate::config::FullDocument as F;
161    match fd {
162        F::Off => None,
163        F::WhenAvailable => Some(FullDocumentType::WhenAvailable),
164        F::Required => Some(FullDocumentType::Required),
165        F::UpdateLookup => Some(FullDocumentType::UpdateLookup),
166    }
167}
168
169/// Map the config pre-image mode to the driver type (`None` = don't request).
170pub(crate) fn full_document_before_type(
171    fd: crate::config::FullDocumentBeforeChange,
172) -> Option<FullDocumentBeforeChangeType> {
173    use crate::config::FullDocumentBeforeChange as F;
174    match fd {
175        F::Off => None,
176        F::WhenAvailable => Some(FullDocumentBeforeChangeType::WhenAvailable),
177        F::Required => Some(FullDocumentBeforeChangeType::Required),
178    }
179}
180
181/// Build the server-side pipeline: operation-type `$match` + user stages.
182pub(crate) fn build_pipeline(config: &MongoCdcSourceConfig) -> Result<Vec<Document>, FaucetError> {
183    let mut pipeline: Vec<Document> = Vec::new();
184    if !config.operation_types.is_empty() {
185        pipeline.push(bson::doc! {
186            "$match": { "operationType": { "$in": config.operation_types.clone() } }
187        });
188    }
189    for (i, stage) in config.aggregation_pipeline.iter().enumerate() {
190        let b = bson::to_bson(stage)
191            .map_err(|e| FaucetError::Config(format!("mongodb-cdc: pipeline stage {i}: {e}")))?;
192        match b {
193            bson::Bson::Document(d) => pipeline.push(d),
194            other => {
195                return Err(FaucetError::Config(format!(
196                    "mongodb-cdc: aggregation_pipeline[{i}] must be an object, got {other:?}"
197                )));
198            }
199        }
200    }
201    Ok(pipeline)
202}
203
204#[async_trait]
205impl Source for MongoCdcSource {
206    async fn fetch_with_context(
207        &self,
208        ctx: &HashMap<String, Value>,
209    ) -> Result<Vec<Value>, FaucetError> {
210        use futures::StreamExt;
211        let mut pages = self.stream_pages(ctx, 0);
212        let mut all = Vec::new();
213        while let Some(page) = pages.next().await {
214            all.extend(page?.records);
215        }
216        Ok(all)
217    }
218
219    fn stream_pages<'a>(
220        &'a self,
221        ctx: &'a HashMap<String, Value>,
222        _batch_size: usize,
223    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
224        self.stream_pages_impl(ctx, self.config.batch_size)
225    }
226
227    fn config_schema(&self) -> Value {
228        serde_json::to_value(schemars::schema_for!(MongoCdcSourceConfig)).unwrap_or(Value::Null)
229    }
230
231    fn state_key(&self) -> Option<String> {
232        Some(self.state_key_value.clone())
233    }
234
235    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
236        let b = Bookmark::from_value(bookmark)?;
237        // Round-trip-validate the token now so a corrupt bookmark fails fast.
238        b.to_token()?;
239        *self.pending_bookmark.lock().await = Some(b);
240        Ok(())
241    }
242
243    async fn capture_resume_position(&self) -> Result<Option<Value>, FaucetError> {
244        let token = self.capture_now_token().await?;
245        Ok(Some(Bookmark::from_token(&token)?.to_value()?))
246    }
247
248    fn supports_exactly_once(&self) -> bool {
249        // Durable resumeToken + deterministic replay from it + per-event
250        // (per-page) bookmarks — the requirements for exactly-once delivery.
251        true
252    }
253
254    fn connector_name(&self) -> &'static str {
255        "mongodb-cdc"
256    }
257
258    fn dataset_uri(&self) -> String {
259        use crate::config::Scope;
260        let base = faucet_core::redact_uri_credentials(&self.config.connection_uri);
261        match &self.config.scope {
262            Scope::Collection {
263                database,
264                collection,
265            } => {
266                format!("{base}/{database}/{collection}")
267            }
268            Scope::Database { database } => format!("{base}/{database}"),
269            Scope::Cluster => base,
270        }
271    }
272
273    async fn check(
274        &self,
275        ctx: &faucet_core::check::CheckContext,
276    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
277        use faucet_core::check::{CheckReport, Probe};
278        let start = std::time::Instant::now();
279
280        // Run `hello` once under the probe timeout, then derive two probes from
281        // the outcome: `connection` (could we reach + auth to the server?) and
282        // `topology` (does it support change streams?). On a connection
283        // failure/timeout the topology probe cannot run, so only `connection`
284        // is reported.
285        let hello = tokio::time::timeout(
286            ctx.timeout,
287            self.client
288                .database("admin")
289                .run_command(bson::doc! { "hello": 1 }),
290        )
291        .await;
292
293        let hello = match hello {
294            Ok(Ok(doc)) => doc,
295            Ok(Err(e)) => {
296                return Ok(CheckReport::single(Probe::fail_hint(
297                    "connection",
298                    start.elapsed(),
299                    format!("could not run hello: {e}"),
300                    "verify the URI is reachable and credentials are valid",
301                )));
302            }
303            Err(_elapsed) => {
304                return Ok(CheckReport::single(Probe::fail_hint(
305                    "connection",
306                    start.elapsed(),
307                    "connection timed out",
308                    "the server did not respond within the check timeout",
309                )));
310            }
311        };
312
313        let connection = Probe::pass("connection", start.elapsed());
314        let topology = if is_changestream_topology(&hello) {
315            Probe::pass("topology", start.elapsed())
316        } else {
317            Probe::fail_hint(
318                "topology",
319                start.elapsed(),
320                "target is a standalone mongod",
321                "Change Streams require a replica set or sharded cluster",
322            )
323        };
324        Ok(CheckReport {
325            probes: vec![connection, topology],
326        })
327    }
328}
329
330impl MongoCdcSource {
331    /// Open the change stream and drain it with an idle-timeout terminator.
332    fn stream_pages_impl<'a>(
333        &'a self,
334        _ctx: &'a HashMap<String, Value>,
335        batch_size: usize,
336    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
337        let idle_timeout = self.config.idle_timeout;
338        let max_await = std::time::Duration::from_millis(self.config.max_await_time_ms);
339        let per_batch = batch_size != 0;
340        let max_staged = self.config.max_staged_records;
341
342        Box::pin(async_stream::try_stream! {
343            use futures::StreamExt;
344
345            let pending = self.pending_bookmark.lock().await.take();
346            let start = resolve_start(&self.config.start_from, pending.as_ref())?;
347            let pipeline = build_pipeline(&self.config)?;
348
349            // Open the change stream scoped to collection, database, or cluster.
350            // Each arm resolves to a ChangeStream<ChangeStreamEvent<Document>>.
351            let change_stream: mongodb::change_stream::ChangeStream<
352                mongodb::change_stream::event::ChangeStreamEvent<Document>,
353            > = match &self.config.scope {
354                crate::config::Scope::Collection { database, collection } => {
355                    let coll: mongodb::Collection<Document> =
356                        self.client.database(database).collection(collection);
357                    let mut w = coll
358                        .watch()
359                        .pipeline(pipeline)
360                        .max_await_time(max_await);
361                    if let Some(fd) = full_document_type(self.config.full_document) {
362                        w = w.full_document(fd);
363                    }
364                    if let Some(fb) =
365                        full_document_before_type(self.config.full_document_before_change)
366                    {
367                        w = w.full_document_before_change(fb);
368                    }
369                    w = match &start {
370                        StartPosition::Now => w,
371                        StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
372                        StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
373                        StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
374                    };
375                    w.await
376                        .map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
377                }
378                crate::config::Scope::Database { database } => {
379                    let db = self.client.database(database);
380                    let mut w = db
381                        .watch()
382                        .pipeline(pipeline)
383                        .max_await_time(max_await);
384                    if let Some(fd) = full_document_type(self.config.full_document) {
385                        w = w.full_document(fd);
386                    }
387                    if let Some(fb) =
388                        full_document_before_type(self.config.full_document_before_change)
389                    {
390                        w = w.full_document_before_change(fb);
391                    }
392                    w = match &start {
393                        StartPosition::Now => w,
394                        StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
395                        StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
396                        StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
397                    };
398                    w.await
399                        .map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
400                }
401                crate::config::Scope::Cluster => {
402                    let mut w = self.client
403                        .watch()
404                        .pipeline(pipeline)
405                        .max_await_time(max_await);
406                    if let Some(fd) = full_document_type(self.config.full_document) {
407                        w = w.full_document(fd);
408                    }
409                    if let Some(fb) =
410                        full_document_before_type(self.config.full_document_before_change)
411                    {
412                        w = w.full_document_before_change(fb);
413                    }
414                    w = match &start {
415                        StartPosition::Now => w,
416                        StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
417                        StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
418                        StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
419                    };
420                    w.await
421                        .map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
422                }
423            };
424            let mut change_stream = std::pin::pin!(change_stream);
425
426            let chunk = if per_batch { batch_size } else { usize::MAX };
427            let mut buffer: Vec<Value> = Vec::new();
428            // Bookmark of the last record currently in `buffer`. `None` only when
429            // the buffer is empty (no events seen this cycle).
430            let mut last_bookmark: Option<Value> = None;
431
432            // Graceful shutdown is handled one layer up: `run_stream`'s page loop
433            // `biased`-`select!`s the pipeline cancel token (Pipeline::with_cancel,
434            // #146) against polling this stream for its next page, so a cancel
435            // mid-idle-wait stops at the page boundary and flushes the sinks — no
436            // source-side signal handling needed. Nothing is persisted until a
437            // page yields, so a dropped in-flight buffer is simply re-fetched
438            // (the resume token has not advanced).
439            loop {
440                match tokio::time::timeout(idle_timeout, change_stream.next()).await {
441                    Ok(Some(Ok(event))) => {
442                        let bookmark = Bookmark::from_token(&event.id)?;
443                        let is_invalidate = matches!(
444                            event.operation_type,
445                            mongodb::change_stream::event::OperationType::Invalidate
446                        );
447                        // OOM safety valve: with `batch_size: 0` (or `fetch_all`)
448                        // the buffer is never flushed mid-cycle, so a
449                        // high-throughput stream would grow it without bound.
450                        // Abort with a typed error before staging one more record
451                        // past the cap rather than risk an OOM-kill.
452                        check_staging_cap(buffer.len(), max_staged)?;
453                        buffer.push(to_envelope(&event, &bookmark)?);
454                        // The persisted bookmark records whether it is an
455                        // invalidate token, so resume uses `start_after` not
456                        // `resume_after` (#321 M3). The envelope above keeps the
457                        // plain token — only the durable state carries the flag.
458                        let persisted = Bookmark {
459                            invalidate: is_invalidate,
460                            ..bookmark
461                        };
462                        last_bookmark = Some(persisted.to_value()?);
463
464                        // An invalidate closes the stream server-side: flush the
465                        // page (including the invalidate event) and end.
466                        if is_invalidate {
467                            yield StreamPage {
468                                records: std::mem::take(&mut buffer),
469                                bookmark: last_bookmark.take(),
470                            };
471                            break;
472                        }
473                        if per_batch && buffer.len() >= chunk {
474                            yield StreamPage {
475                                records: std::mem::take(&mut buffer),
476                                bookmark: last_bookmark.take(),
477                            };
478                        }
479                    }
480                    Ok(Some(Err(e))) => {
481                        Err(FaucetError::Source(format!("mongodb-cdc stream error: {e}")))?;
482                    }
483                    // Idle (timeout) or cursor closed: flush whatever we have and
484                    // end this fetch cycle.
485                    Ok(None) | Err(_) => {
486                        if !buffer.is_empty() {
487                            yield StreamPage {
488                                records: std::mem::take(&mut buffer),
489                                bookmark: last_bookmark.take(),
490                            };
491                        }
492                        break;
493                    }
494                }
495            }
496
497            tracing::info!(connector = "mongodb-cdc", "change stream fetch cycle complete");
498        })
499    }
500
501    /// Open a change stream at "now", read its `postBatchResumeToken` without
502    /// consuming any events, then drop the stream. Used by the replication
503    /// snapshot→CDC handoff to anchor the stream before the bulk snapshot.
504    async fn capture_now_token(&self) -> Result<ResumeToken, FaucetError> {
505        use crate::config::Scope;
506        use futures::StreamExt;
507        let max_await = std::time::Duration::from_millis(self.config.max_await_time_ms);
508        let cs: mongodb::change_stream::ChangeStream<
509            mongodb::change_stream::event::ChangeStreamEvent<Document>,
510        > = match &self.config.scope {
511            Scope::Collection {
512                database,
513                collection,
514            } => self
515                .client
516                .database(database)
517                .collection::<Document>(collection)
518                .watch()
519                .max_await_time(max_await)
520                .await
521                .map_err(|e| {
522                    FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
523                })?,
524            Scope::Database { database } => self
525                .client
526                .database(database)
527                .watch()
528                .max_await_time(max_await)
529                .await
530                .map_err(|e| {
531                    FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
532                })?,
533            Scope::Cluster => self
534                .client
535                .watch()
536                .max_await_time(max_await)
537                .await
538                .map_err(|e| {
539                    FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
540                })?,
541        };
542        let mut cs = std::pin::pin!(cs);
543        // The initial aggregate response usually carries a postBatchResumeToken.
544        // If the driver hasn't cached one yet, one short poll forces a getMore
545        // that returns it.
546        if let Some(tok) = cs.resume_token() {
547            return Ok(tok);
548        }
549        let _ = tokio::time::timeout(
550            max_await.max(std::time::Duration::from_millis(500)),
551            cs.next(),
552        )
553        .await;
554        cs.resume_token().ok_or_else(|| {
555            FaucetError::Source(
556                "mongodb-cdc: could not capture a resume token (no postBatchResumeToken returned)"
557                    .into(),
558            )
559        })
560    }
561}
562
563#[cfg(test)]
564mod tests {
565    use super::*;
566    use serde_json::json;
567
568    #[test]
569    fn explicit_timestamp_overrides_bookmark() {
570        let pending = Bookmark {
571            resume_token: json!({ "_data": "AA" }),
572            ..Default::default()
573        };
574        let pos = resolve_start(
575            &StartFrom::Timestamp {
576                timestamp_secs: 100,
577            },
578            Some(&pending),
579        )
580        .unwrap();
581        assert_eq!(
582            pos,
583            StartPosition::AtOperationTime(Timestamp {
584                time: 100,
585                increment: 0
586            })
587        );
588    }
589
590    #[test]
591    fn explicit_resume_token_overrides_bookmark() {
592        // An explicit ResumeToken in config must win over any persisted
593        // bookmark and resolve to ResumeAfter(token).
594        let pending = Bookmark {
595            resume_token: json!({ "_data": "PENDING" }),
596            ..Default::default()
597        };
598        let pos = resolve_start(
599            &StartFrom::ResumeToken {
600                token: json!({ "_data": "8264AB00" }),
601            },
602            Some(&pending),
603        )
604        .unwrap();
605        // The resolved position must round-trip back to the configured token,
606        // not the pending bookmark's token.
607        match pos {
608            StartPosition::ResumeAfter(tok) => {
609                let b = Bookmark::from_token(&tok).unwrap();
610                assert_eq!(b.resume_token["_data"], json!("8264AB00"));
611            }
612            other => panic!("expected ResumeAfter, got {other:?}"),
613        }
614    }
615
616    #[test]
617    fn now_yields_to_bookmark() {
618        let pending = Bookmark {
619            resume_token: json!({ "_data": "8264AB00" }),
620            ..Default::default()
621        };
622        let pos = resolve_start(&StartFrom::Now, Some(&pending)).unwrap();
623        assert!(matches!(pos, StartPosition::ResumeAfter(_)));
624    }
625
626    #[test]
627    fn invalidate_bookmark_resumes_with_start_after() {
628        // #321 M3: a bookmark captured from an invalidate event must resolve to
629        // StartAfter (MongoDB rejects resumeAfter on an invalidate token).
630        let pending = Bookmark {
631            resume_token: json!({ "_data": "8264AB00" }),
632            invalidate: true,
633        };
634        let pos = resolve_start(&StartFrom::Now, Some(&pending)).unwrap();
635        assert!(
636            matches!(pos, StartPosition::StartAfter(_)),
637            "invalidate bookmark must use start_after, got {pos:?}"
638        );
639        // A non-invalidate bookmark still uses resume_after.
640        let normal = Bookmark {
641            resume_token: json!({ "_data": "8264AB00" }),
642            invalidate: false,
643        };
644        assert!(matches!(
645            resolve_start(&StartFrom::Earliest, Some(&normal)).unwrap(),
646            StartPosition::ResumeAfter(_)
647        ));
648    }
649
650    #[test]
651    fn now_without_bookmark_is_now() {
652        assert_eq!(
653            resolve_start(&StartFrom::Now, None).unwrap(),
654            StartPosition::Now
655        );
656    }
657
658    #[test]
659    fn earliest_without_bookmark_uses_operation_time() {
660        assert_eq!(
661            resolve_start(&StartFrom::Earliest, None).unwrap(),
662            StartPosition::AtOperationTime(Timestamp {
663                time: 1,
664                increment: 1
665            })
666        );
667    }
668
669    #[test]
670    fn staging_cap_none_is_unbounded() {
671        // With no cap, an arbitrarily large buffer never aborts.
672        assert!(check_staging_cap(0, None).is_ok());
673        assert!(check_staging_cap(1_000_000, None).is_ok());
674    }
675
676    #[test]
677    fn staging_cap_allows_up_to_limit() {
678        // Buffering record N is allowed while the current length is below the
679        // cap (i.e. up to `max` records total).
680        assert!(check_staging_cap(0, Some(2)).is_ok());
681        assert!(check_staging_cap(1, Some(2)).is_ok());
682    }
683
684    #[test]
685    fn staging_cap_aborts_at_limit() {
686        // Once `current_len` reaches the cap, staging one more record must
687        // abort with a typed FaucetError::Source naming max_staged_records.
688        let err = check_staging_cap(2, Some(2)).unwrap_err();
689        assert!(matches!(err, FaucetError::Source(_)));
690        assert!(format!("{err}").contains("max_staged_records"));
691
692        // And it stays an error for any length beyond the cap.
693        assert!(check_staging_cap(3, Some(2)).is_err());
694    }
695
696    #[test]
697    fn staging_cap_drives_accumulation_past_limit() {
698        // Simulate the buffer-accumulation loop with batch_size: 0 (no
699        // mid-cycle flush) and a cap of 3: the 4th record must abort.
700        let mut buffer: Vec<u8> = Vec::new();
701        let max = Some(3usize);
702        let mut last: Result<(), FaucetError> = Ok(());
703        for i in 0..10u8 {
704            if let Err(e) = check_staging_cap(buffer.len(), max) {
705                last = Err(e);
706                break;
707            }
708            buffer.push(i);
709        }
710        assert_eq!(buffer.len(), 3, "buffer stops growing at the cap");
711        let err = last.unwrap_err();
712        assert!(matches!(err, FaucetError::Source(_)));
713        assert!(format!("{err}").contains("max_staged_records"));
714    }
715
716    #[test]
717    fn topology_classification() {
718        assert!(is_changestream_topology(&bson::doc! { "setName": "rs0" }));
719        assert!(is_changestream_topology(&bson::doc! { "msg": "isdbgrid" }));
720        assert!(!is_changestream_topology(&bson::doc! { "ok": 1.0 }));
721    }
722
723    #[test]
724    fn pipeline_prepends_operation_match() {
725        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
726            "connection_uri": "mongodb://h/?replicaSet=rs0",
727            "scope": { "type": "cluster" },
728            "operation_types": ["insert", "delete"]
729        }))
730        .unwrap();
731        let p = build_pipeline(&c).unwrap();
732        assert_eq!(p.len(), 1);
733        assert!(p[0].contains_key("$match"));
734    }
735
736    #[test]
737    fn pipeline_rejects_non_object_stage() {
738        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
739            "connection_uri": "mongodb://h/?replicaSet=rs0",
740            "scope": { "type": "cluster" },
741            "aggregation_pipeline": [[1, 2, 3]]
742        }))
743        .unwrap();
744        assert!(matches!(build_pipeline(&c), Err(FaucetError::Config(_))));
745    }
746
747    #[test]
748    fn pipeline_appends_object_stages_after_match() {
749        // operation_types prepends a $match; each aggregation_pipeline object
750        // stage is appended in order.
751        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
752            "connection_uri": "mongodb://h/?replicaSet=rs0",
753            "scope": { "type": "cluster" },
754            "operation_types": ["insert"],
755            "aggregation_pipeline": [
756                { "$match": { "fullDocument.tier": "gold" } },
757                { "$project": { "fullDocument._id": 1 } }
758            ]
759        }))
760        .unwrap();
761        let p = build_pipeline(&c).unwrap();
762        assert_eq!(p.len(), 3, "operation $match + two user stages");
763        // First is the operation-type $match.
764        assert!(
765            p[0].get_document("$match")
766                .unwrap()
767                .contains_key("operationType")
768        );
769        // User object stages preserved verbatim, in order.
770        assert!(p[1].contains_key("$match"));
771        assert!(p[2].contains_key("$project"));
772    }
773
774    #[test]
775    fn pipeline_object_stages_without_operation_match() {
776        // No operation_types → no leading $match; only the user object stage.
777        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
778            "connection_uri": "mongodb://h/?replicaSet=rs0",
779            "scope": { "type": "cluster" },
780            "aggregation_pipeline": [ { "$project": { "_id": 1 } } ]
781        }))
782        .unwrap();
783        let p = build_pipeline(&c).unwrap();
784        assert_eq!(p.len(), 1);
785        assert!(p[0].contains_key("$project"));
786    }
787
788    #[test]
789    fn full_document_type_mapping() {
790        use crate::config::FullDocument as F;
791        assert!(full_document_type(F::Off).is_none());
792        assert!(matches!(
793            full_document_type(F::WhenAvailable),
794            Some(FullDocumentType::WhenAvailable)
795        ));
796        assert!(matches!(
797            full_document_type(F::Required),
798            Some(FullDocumentType::Required)
799        ));
800        assert!(matches!(
801            full_document_type(F::UpdateLookup),
802            Some(FullDocumentType::UpdateLookup)
803        ));
804    }
805
806    #[test]
807    fn full_document_before_type_mapping() {
808        use crate::config::FullDocumentBeforeChange as F;
809        assert!(full_document_before_type(F::Off).is_none());
810        assert!(matches!(
811            full_document_before_type(F::WhenAvailable),
812            Some(FullDocumentBeforeChangeType::WhenAvailable)
813        ));
814        assert!(matches!(
815            full_document_before_type(F::Required),
816            Some(FullDocumentBeforeChangeType::Required)
817        ));
818    }
819
820    // dataset_uri is a pure-config method; MongoCdcSource requires a live
821    // MongoDB async constructor, so we verify the logic directly using config
822    // deserialization (which is sync).
823    #[test]
824    fn dataset_uri_cluster_scope() {
825        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
826            "connection_uri": "mongodb://u:p@h:27017/?replicaSet=rs0",
827            "scope": { "type": "cluster" }
828        }))
829        .unwrap();
830        use crate::config::Scope;
831        let base = faucet_core::redact_uri_credentials(&c.connection_uri);
832        let uri = match &c.scope {
833            Scope::Collection {
834                database,
835                collection,
836            } => format!("{base}/{database}/{collection}"),
837            Scope::Database { database } => format!("{base}/{database}"),
838            Scope::Cluster => base,
839        };
840        assert_eq!(uri, "mongodb://h:27017/?replicaSet=rs0");
841    }
842
843    #[test]
844    fn dataset_uri_collection_scope() {
845        let c: MongoCdcSourceConfig = serde_json::from_value(json!({
846            "connection_uri": "mongodb://u:p@h:27017/?replicaSet=rs0",
847            "scope": { "type": "collection", "database": "mydb", "collection": "orders" }
848        }))
849        .unwrap();
850        use crate::config::Scope;
851        let base = faucet_core::redact_uri_credentials(&c.connection_uri);
852        let uri = match &c.scope {
853            Scope::Collection {
854                database,
855                collection,
856            } => format!("{base}/{database}/{collection}"),
857            Scope::Database { database } => format!("{base}/{database}"),
858            Scope::Cluster => base,
859        };
860        assert_eq!(uri, "mongodb://h:27017/?replicaSet=rs0/mydb/orders");
861    }
862}