use crate::config::{MongoCdcSourceConfig, StartFrom};
use crate::envelope::to_envelope;
use crate::state::{Bookmark, state_key};
use async_trait::async_trait;
use faucet_core::{FaucetError, Source, Stream, StreamPage};
use mongodb::Client;
use mongodb::bson::{self, Document, Timestamp};
use mongodb::change_stream::event::ResumeToken;
use mongodb::options::{FullDocumentBeforeChangeType, FullDocumentType};
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;
use tokio::sync::Mutex;
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum StartPosition {
Now,
AtOperationTime(Timestamp),
ResumeAfter(ResumeToken),
StartAfter(ResumeToken),
}
pub(crate) fn resolve_start(
start_from: &StartFrom,
pending: Option<&Bookmark>,
) -> Result<StartPosition, FaucetError> {
match start_from {
StartFrom::ResumeToken { token } => {
let b = Bookmark {
resume_token: token.clone(),
..Default::default()
};
Ok(StartPosition::ResumeAfter(b.to_token()?))
}
StartFrom::Timestamp { timestamp_secs } => Ok(StartPosition::AtOperationTime(Timestamp {
time: *timestamp_secs,
increment: 0,
})),
StartFrom::Now | StartFrom::Earliest => {
if let Some(b) = pending {
if b.invalidate {
Ok(StartPosition::StartAfter(b.to_token()?))
} else {
Ok(StartPosition::ResumeAfter(b.to_token()?))
}
} else if matches!(start_from, StartFrom::Earliest) {
Ok(StartPosition::AtOperationTime(Timestamp {
time: 1,
increment: 1,
}))
} else {
Ok(StartPosition::Now)
}
}
}
}
pub(crate) fn check_staging_cap(
current_len: usize,
max_staged: Option<usize>,
) -> Result<(), FaucetError> {
if let Some(max) = max_staged
&& current_len >= max
{
return Err(FaucetError::Source(format!(
"mongodb-cdc: in-memory change buffer exceeded max_staged_records ({max}); \
aborting to avoid unbounded memory growth. Raise max_staged_records or \
reduce batch_size so pages flush more often."
)));
}
Ok(())
}
pub struct MongoCdcSource {
config: MongoCdcSourceConfig,
client: Client,
state_key_value: String,
pending_bookmark: Mutex<Option<Bookmark>>,
}
impl MongoCdcSource {
pub async fn new(config: MongoCdcSourceConfig) -> Result<Self, FaucetError> {
config.validate()?;
let client = Client::with_uri_str(&config.connection_uri)
.await
.map_err(|e| FaucetError::Source(format!("MongoDB connection failed: {e}")))?;
ensure_changestream_capable(&client).await?;
if config.full_document_before_change != crate::config::FullDocumentBeforeChange::Off {
tracing::warn!(
"mongodb-cdc: full_document_before_change requires MongoDB 6.0+ and the \
collection's changeStreamPreAndPostImages enabled; the server will error \
at stream open if unavailable"
);
}
let state_key_value = state_key(&config.scope);
Ok(Self {
config,
client,
state_key_value,
pending_bookmark: Mutex::new(None),
})
}
}
async fn ensure_changestream_capable(client: &Client) -> Result<(), FaucetError> {
let hello = client
.database("admin")
.run_command(bson::doc! { "hello": 1 })
.await
.map_err(|e| FaucetError::Source(format!("mongodb-cdc hello failed: {e}")))?;
if is_changestream_topology(&hello) {
Ok(())
} else {
Err(FaucetError::Source(
"mongodb-cdc: Change Streams require a replica set or sharded cluster; \
the target appears to be a standalone mongod"
.into(),
))
}
}
pub(crate) fn is_changestream_topology(hello: &Document) -> bool {
hello.get_str("setName").is_ok()
|| hello
.get_str("msg")
.map(|m| m == "isdbgrid")
.unwrap_or(false)
}
pub(crate) fn full_document_type(fd: crate::config::FullDocument) -> Option<FullDocumentType> {
use crate::config::FullDocument as F;
match fd {
F::Off => None,
F::WhenAvailable => Some(FullDocumentType::WhenAvailable),
F::Required => Some(FullDocumentType::Required),
F::UpdateLookup => Some(FullDocumentType::UpdateLookup),
}
}
pub(crate) fn full_document_before_type(
fd: crate::config::FullDocumentBeforeChange,
) -> Option<FullDocumentBeforeChangeType> {
use crate::config::FullDocumentBeforeChange as F;
match fd {
F::Off => None,
F::WhenAvailable => Some(FullDocumentBeforeChangeType::WhenAvailable),
F::Required => Some(FullDocumentBeforeChangeType::Required),
}
}
pub(crate) fn build_pipeline(config: &MongoCdcSourceConfig) -> Result<Vec<Document>, FaucetError> {
let mut pipeline: Vec<Document> = Vec::new();
if !config.operation_types.is_empty() {
pipeline.push(bson::doc! {
"$match": { "operationType": { "$in": config.operation_types.clone() } }
});
}
for (i, stage) in config.aggregation_pipeline.iter().enumerate() {
let b = bson::to_bson(stage)
.map_err(|e| FaucetError::Config(format!("mongodb-cdc: pipeline stage {i}: {e}")))?;
match b {
bson::Bson::Document(d) => pipeline.push(d),
other => {
return Err(FaucetError::Config(format!(
"mongodb-cdc: aggregation_pipeline[{i}] must be an object, got {other:?}"
)));
}
}
}
Ok(pipeline)
}
#[async_trait]
impl Source for MongoCdcSource {
async fn fetch_with_context(
&self,
ctx: &HashMap<String, Value>,
) -> Result<Vec<Value>, FaucetError> {
use futures::StreamExt;
let mut pages = self.stream_pages(ctx, 0);
let mut all = Vec::new();
while let Some(page) = pages.next().await {
all.extend(page?.records);
}
Ok(all)
}
fn stream_pages<'a>(
&'a self,
ctx: &'a HashMap<String, Value>,
_batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
self.stream_pages_impl(ctx, self.config.batch_size)
}
fn config_schema(&self) -> Value {
serde_json::to_value(schemars::schema_for!(MongoCdcSourceConfig)).unwrap_or(Value::Null)
}
fn state_key(&self) -> Option<String> {
Some(self.state_key_value.clone())
}
async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
let b = Bookmark::from_value(bookmark)?;
b.to_token()?;
*self.pending_bookmark.lock().await = Some(b);
Ok(())
}
async fn capture_resume_position(&self) -> Result<Option<Value>, FaucetError> {
let token = self.capture_now_token().await?;
Ok(Some(Bookmark::from_token(&token)?.to_value()?))
}
fn supports_exactly_once(&self) -> bool {
true
}
fn connector_name(&self) -> &'static str {
"mongodb-cdc"
}
fn dataset_uri(&self) -> String {
use crate::config::Scope;
let base = faucet_core::redact_uri_credentials(&self.config.connection_uri);
match &self.config.scope {
Scope::Collection {
database,
collection,
} => {
format!("{base}/{database}/{collection}")
}
Scope::Database { database } => format!("{base}/{database}"),
Scope::Cluster => base,
}
}
async fn check(
&self,
ctx: &faucet_core::check::CheckContext,
) -> Result<faucet_core::check::CheckReport, FaucetError> {
use faucet_core::check::{CheckReport, Probe};
let start = std::time::Instant::now();
let hello = tokio::time::timeout(
ctx.timeout,
self.client
.database("admin")
.run_command(bson::doc! { "hello": 1 }),
)
.await;
let hello = match hello {
Ok(Ok(doc)) => doc,
Ok(Err(e)) => {
return Ok(CheckReport::single(Probe::fail_hint(
"connection",
start.elapsed(),
format!("could not run hello: {e}"),
"verify the URI is reachable and credentials are valid",
)));
}
Err(_elapsed) => {
return Ok(CheckReport::single(Probe::fail_hint(
"connection",
start.elapsed(),
"connection timed out",
"the server did not respond within the check timeout",
)));
}
};
let connection = Probe::pass("connection", start.elapsed());
let topology = if is_changestream_topology(&hello) {
Probe::pass("topology", start.elapsed())
} else {
Probe::fail_hint(
"topology",
start.elapsed(),
"target is a standalone mongod",
"Change Streams require a replica set or sharded cluster",
)
};
Ok(CheckReport {
probes: vec![connection, topology],
})
}
}
impl MongoCdcSource {
fn stream_pages_impl<'a>(
&'a self,
_ctx: &'a HashMap<String, Value>,
batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
let idle_timeout = self.config.idle_timeout;
let max_await = std::time::Duration::from_millis(self.config.max_await_time_ms);
let per_batch = batch_size != 0;
let max_staged = self.config.max_staged_records;
Box::pin(async_stream::try_stream! {
use futures::StreamExt;
let pending = self.pending_bookmark.lock().await.take();
let start = resolve_start(&self.config.start_from, pending.as_ref())?;
let pipeline = build_pipeline(&self.config)?;
let change_stream: mongodb::change_stream::ChangeStream<
mongodb::change_stream::event::ChangeStreamEvent<Document>,
> = match &self.config.scope {
crate::config::Scope::Collection { database, collection } => {
let coll: mongodb::Collection<Document> =
self.client.database(database).collection(collection);
let mut w = coll
.watch()
.pipeline(pipeline)
.max_await_time(max_await);
if let Some(fd) = full_document_type(self.config.full_document) {
w = w.full_document(fd);
}
if let Some(fb) =
full_document_before_type(self.config.full_document_before_change)
{
w = w.full_document_before_change(fb);
}
w = match &start {
StartPosition::Now => w,
StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
};
w.await
.map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
}
crate::config::Scope::Database { database } => {
let db = self.client.database(database);
let mut w = db
.watch()
.pipeline(pipeline)
.max_await_time(max_await);
if let Some(fd) = full_document_type(self.config.full_document) {
w = w.full_document(fd);
}
if let Some(fb) =
full_document_before_type(self.config.full_document_before_change)
{
w = w.full_document_before_change(fb);
}
w = match &start {
StartPosition::Now => w,
StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
};
w.await
.map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
}
crate::config::Scope::Cluster => {
let mut w = self.client
.watch()
.pipeline(pipeline)
.max_await_time(max_await);
if let Some(fd) = full_document_type(self.config.full_document) {
w = w.full_document(fd);
}
if let Some(fb) =
full_document_before_type(self.config.full_document_before_change)
{
w = w.full_document_before_change(fb);
}
w = match &start {
StartPosition::Now => w,
StartPosition::AtOperationTime(ts) => w.start_at_operation_time(*ts),
StartPosition::ResumeAfter(tok) => w.resume_after(tok.clone()),
StartPosition::StartAfter(tok) => w.start_after(tok.clone()),
};
w.await
.map_err(|e| FaucetError::Source(format!("mongodb-cdc watch failed: {e}")))?
}
};
let mut change_stream = std::pin::pin!(change_stream);
let chunk = if per_batch { batch_size } else { usize::MAX };
let mut buffer: Vec<Value> = Vec::new();
let mut last_bookmark: Option<Value> = None;
loop {
match tokio::time::timeout(idle_timeout, change_stream.next()).await {
Ok(Some(Ok(event))) => {
let bookmark = Bookmark::from_token(&event.id)?;
let is_invalidate = matches!(
event.operation_type,
mongodb::change_stream::event::OperationType::Invalidate
);
check_staging_cap(buffer.len(), max_staged)?;
buffer.push(to_envelope(&event, &bookmark)?);
let persisted = Bookmark {
invalidate: is_invalidate,
..bookmark
};
last_bookmark = Some(persisted.to_value()?);
if is_invalidate {
yield StreamPage {
records: std::mem::take(&mut buffer),
bookmark: last_bookmark.take(),
};
break;
}
if per_batch && buffer.len() >= chunk {
yield StreamPage {
records: std::mem::take(&mut buffer),
bookmark: last_bookmark.take(),
};
}
}
Ok(Some(Err(e))) => {
Err(FaucetError::Source(format!("mongodb-cdc stream error: {e}")))?;
}
Ok(None) | Err(_) => {
if !buffer.is_empty() {
yield StreamPage {
records: std::mem::take(&mut buffer),
bookmark: last_bookmark.take(),
};
}
break;
}
}
}
tracing::info!(connector = "mongodb-cdc", "change stream fetch cycle complete");
})
}
async fn capture_now_token(&self) -> Result<ResumeToken, FaucetError> {
use crate::config::Scope;
use futures::StreamExt;
let max_await = std::time::Duration::from_millis(self.config.max_await_time_ms);
let cs: mongodb::change_stream::ChangeStream<
mongodb::change_stream::event::ChangeStreamEvent<Document>,
> = match &self.config.scope {
Scope::Collection {
database,
collection,
} => self
.client
.database(database)
.collection::<Document>(collection)
.watch()
.max_await_time(max_await)
.await
.map_err(|e| {
FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
})?,
Scope::Database { database } => self
.client
.database(database)
.watch()
.max_await_time(max_await)
.await
.map_err(|e| {
FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
})?,
Scope::Cluster => self
.client
.watch()
.max_await_time(max_await)
.await
.map_err(|e| {
FaucetError::Source(format!("mongodb-cdc capture watch failed: {e}"))
})?,
};
let mut cs = std::pin::pin!(cs);
if let Some(tok) = cs.resume_token() {
return Ok(tok);
}
let _ = tokio::time::timeout(
max_await.max(std::time::Duration::from_millis(500)),
cs.next(),
)
.await;
cs.resume_token().ok_or_else(|| {
FaucetError::Source(
"mongodb-cdc: could not capture a resume token (no postBatchResumeToken returned)"
.into(),
)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn explicit_timestamp_overrides_bookmark() {
let pending = Bookmark {
resume_token: json!({ "_data": "AA" }),
..Default::default()
};
let pos = resolve_start(
&StartFrom::Timestamp {
timestamp_secs: 100,
},
Some(&pending),
)
.unwrap();
assert_eq!(
pos,
StartPosition::AtOperationTime(Timestamp {
time: 100,
increment: 0
})
);
}
#[test]
fn explicit_resume_token_overrides_bookmark() {
let pending = Bookmark {
resume_token: json!({ "_data": "PENDING" }),
..Default::default()
};
let pos = resolve_start(
&StartFrom::ResumeToken {
token: json!({ "_data": "8264AB00" }),
},
Some(&pending),
)
.unwrap();
match pos {
StartPosition::ResumeAfter(tok) => {
let b = Bookmark::from_token(&tok).unwrap();
assert_eq!(b.resume_token["_data"], json!("8264AB00"));
}
other => panic!("expected ResumeAfter, got {other:?}"),
}
}
#[test]
fn now_yields_to_bookmark() {
let pending = Bookmark {
resume_token: json!({ "_data": "8264AB00" }),
..Default::default()
};
let pos = resolve_start(&StartFrom::Now, Some(&pending)).unwrap();
assert!(matches!(pos, StartPosition::ResumeAfter(_)));
}
#[test]
fn invalidate_bookmark_resumes_with_start_after() {
let pending = Bookmark {
resume_token: json!({ "_data": "8264AB00" }),
invalidate: true,
};
let pos = resolve_start(&StartFrom::Now, Some(&pending)).unwrap();
assert!(
matches!(pos, StartPosition::StartAfter(_)),
"invalidate bookmark must use start_after, got {pos:?}"
);
let normal = Bookmark {
resume_token: json!({ "_data": "8264AB00" }),
invalidate: false,
};
assert!(matches!(
resolve_start(&StartFrom::Earliest, Some(&normal)).unwrap(),
StartPosition::ResumeAfter(_)
));
}
#[test]
fn now_without_bookmark_is_now() {
assert_eq!(
resolve_start(&StartFrom::Now, None).unwrap(),
StartPosition::Now
);
}
#[test]
fn earliest_without_bookmark_uses_operation_time() {
assert_eq!(
resolve_start(&StartFrom::Earliest, None).unwrap(),
StartPosition::AtOperationTime(Timestamp {
time: 1,
increment: 1
})
);
}
#[test]
fn staging_cap_none_is_unbounded() {
assert!(check_staging_cap(0, None).is_ok());
assert!(check_staging_cap(1_000_000, None).is_ok());
}
#[test]
fn staging_cap_allows_up_to_limit() {
assert!(check_staging_cap(0, Some(2)).is_ok());
assert!(check_staging_cap(1, Some(2)).is_ok());
}
#[test]
fn staging_cap_aborts_at_limit() {
let err = check_staging_cap(2, Some(2)).unwrap_err();
assert!(matches!(err, FaucetError::Source(_)));
assert!(format!("{err}").contains("max_staged_records"));
assert!(check_staging_cap(3, Some(2)).is_err());
}
#[test]
fn staging_cap_drives_accumulation_past_limit() {
let mut buffer: Vec<u8> = Vec::new();
let max = Some(3usize);
let mut last: Result<(), FaucetError> = Ok(());
for i in 0..10u8 {
if let Err(e) = check_staging_cap(buffer.len(), max) {
last = Err(e);
break;
}
buffer.push(i);
}
assert_eq!(buffer.len(), 3, "buffer stops growing at the cap");
let err = last.unwrap_err();
assert!(matches!(err, FaucetError::Source(_)));
assert!(format!("{err}").contains("max_staged_records"));
}
#[test]
fn topology_classification() {
assert!(is_changestream_topology(&bson::doc! { "setName": "rs0" }));
assert!(is_changestream_topology(&bson::doc! { "msg": "isdbgrid" }));
assert!(!is_changestream_topology(&bson::doc! { "ok": 1.0 }));
}
#[test]
fn pipeline_prepends_operation_match() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://h/?replicaSet=rs0",
"scope": { "type": "cluster" },
"operation_types": ["insert", "delete"]
}))
.unwrap();
let p = build_pipeline(&c).unwrap();
assert_eq!(p.len(), 1);
assert!(p[0].contains_key("$match"));
}
#[test]
fn pipeline_rejects_non_object_stage() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://h/?replicaSet=rs0",
"scope": { "type": "cluster" },
"aggregation_pipeline": [[1, 2, 3]]
}))
.unwrap();
assert!(matches!(build_pipeline(&c), Err(FaucetError::Config(_))));
}
#[test]
fn pipeline_appends_object_stages_after_match() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://h/?replicaSet=rs0",
"scope": { "type": "cluster" },
"operation_types": ["insert"],
"aggregation_pipeline": [
{ "$match": { "fullDocument.tier": "gold" } },
{ "$project": { "fullDocument._id": 1 } }
]
}))
.unwrap();
let p = build_pipeline(&c).unwrap();
assert_eq!(p.len(), 3, "operation $match + two user stages");
assert!(
p[0].get_document("$match")
.unwrap()
.contains_key("operationType")
);
assert!(p[1].contains_key("$match"));
assert!(p[2].contains_key("$project"));
}
#[test]
fn pipeline_object_stages_without_operation_match() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://h/?replicaSet=rs0",
"scope": { "type": "cluster" },
"aggregation_pipeline": [ { "$project": { "_id": 1 } } ]
}))
.unwrap();
let p = build_pipeline(&c).unwrap();
assert_eq!(p.len(), 1);
assert!(p[0].contains_key("$project"));
}
#[test]
fn full_document_type_mapping() {
use crate::config::FullDocument as F;
assert!(full_document_type(F::Off).is_none());
assert!(matches!(
full_document_type(F::WhenAvailable),
Some(FullDocumentType::WhenAvailable)
));
assert!(matches!(
full_document_type(F::Required),
Some(FullDocumentType::Required)
));
assert!(matches!(
full_document_type(F::UpdateLookup),
Some(FullDocumentType::UpdateLookup)
));
}
#[test]
fn full_document_before_type_mapping() {
use crate::config::FullDocumentBeforeChange as F;
assert!(full_document_before_type(F::Off).is_none());
assert!(matches!(
full_document_before_type(F::WhenAvailable),
Some(FullDocumentBeforeChangeType::WhenAvailable)
));
assert!(matches!(
full_document_before_type(F::Required),
Some(FullDocumentBeforeChangeType::Required)
));
}
#[test]
fn dataset_uri_cluster_scope() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://u:p@h:27017/?replicaSet=rs0",
"scope": { "type": "cluster" }
}))
.unwrap();
use crate::config::Scope;
let base = faucet_core::redact_uri_credentials(&c.connection_uri);
let uri = match &c.scope {
Scope::Collection {
database,
collection,
} => format!("{base}/{database}/{collection}"),
Scope::Database { database } => format!("{base}/{database}"),
Scope::Cluster => base,
};
assert_eq!(uri, "mongodb://h:27017/?replicaSet=rs0");
}
#[test]
fn dataset_uri_collection_scope() {
let c: MongoCdcSourceConfig = serde_json::from_value(json!({
"connection_uri": "mongodb://u:p@h:27017/?replicaSet=rs0",
"scope": { "type": "collection", "database": "mydb", "collection": "orders" }
}))
.unwrap();
use crate::config::Scope;
let base = faucet_core::redact_uri_credentials(&c.connection_uri);
let uri = match &c.scope {
Scope::Collection {
database,
collection,
} => format!("{base}/{database}/{collection}"),
Scope::Database { database } => format!("{base}/{database}"),
Scope::Cluster => base,
};
assert_eq!(uri, "mongodb://h:27017/?replicaSet=rs0/mydb/orders");
}
}