laminar_connectors/mongodb/source/
lifecycle.rs

1//! Source contract, lifecycle, polling, and shutdown.
2
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use arrow_schema::SchemaRef;
7use async_trait::async_trait;
8use tokio::sync::Notify;
9
10use crate::checkpoint::SourceCheckpoint;
11use crate::config::{ConnectorConfig, ConnectorState};
12use crate::connector::{
13    ConnectorTaskTracker, SourceBatch, SourceConnector, SourceConsistency, SourceContract,
14    SourceInputMode, SourcePosition, SourceRowPositionCapability, SourceStart, SourceTopology,
15};
16use crate::error::ConnectorError;
17
18use super::super::config::{SnapshotMode, SourceOutputMode};
19use super::admission::ReaderLaunch;
20use super::checkpoint::set_emitted_offsets;
21use super::{
22    mongodb_stream_identity, parse_mongodb_checkpoint, reap_mongo_reader, DocumentProjection,
23    MongoCheckpointPosition, MongoDbCdcSource, MongoDbSourceConfig, MongoResumePosition,
24    ParsedMongoCheckpoint, ReaderStart, StreamPosition, COLLECTION_UUID_METADATA,
25    DEPLOYMENT_IDENTITY_METADATA, MONGODB_CHECKPOINT_CONNECTOR, MONGODB_CHECKPOINT_VERSION,
26    READER_SHUTDOWN_TIMEOUT, STREAM_IDENTITY_METADATA,
27};
28
29fn parsed_config(
30    source: &MongoDbSourceConfig,
31    config: &ConnectorConfig,
32) -> Result<MongoDbSourceConfig, ConnectorError> {
33    if config.properties().is_empty() {
34        let mut config = source.clone();
35        config.normalize_pipeline()?;
36        config.validate()?;
37        Ok(config)
38    } else {
39        MongoDbSourceConfig::from_config(config)
40    }
41}
42
43pub(super) fn declared_primary_key(config: &ConnectorConfig) -> Vec<String> {
44    config
45        .get("_primary_key_columns")
46        .unwrap_or("")
47        .split(',')
48        .map(str::trim)
49        .filter(|column| !column.is_empty())
50        .map(str::to_string)
51        .collect()
52}
53
54/// Document-mode projection from the engine-supplied declared schema and primary key.
55pub(super) fn document_projection(
56    parsed: &MongoDbSourceConfig,
57    config: &ConnectorConfig,
58    schema: Option<SchemaRef>,
59) -> Result<Option<DocumentProjection>, ConnectorError> {
60    if parsed.output_mode != SourceOutputMode::Document {
61        return Ok(None);
62    }
63    let schema = schema.ok_or_else(|| {
64        ConnectorError::ConfigurationError(
65            "MongoDB CDC output.mode=document requires declared columns and a PRIMARY KEY".into(),
66        )
67    })?;
68    DocumentProjection::try_new(
69        &schema,
70        &declared_primary_key(config),
71        &parsed.objectid_columns,
72        parsed.document_json_column.as_deref(),
73    )
74    .map(Some)
75}
76
77fn restored_start(checkpoint: &ParsedMongoCheckpoint) -> Result<ReaderStart, ConnectorError> {
78    let token = |token: &str| {
79        serde_json::from_str(token).map_err(|error| {
80            ConnectorError::ConfigurationError(format!(
81                "invalid MongoDB CDC resume token in checkpoint: {error}"
82            ))
83        })
84    };
85    Ok(match &checkpoint.emitted.position {
86        MongoCheckpointPosition::Stream(StreamPosition::ResumeAfter(encoded)) => {
87            ReaderStart::Stream(MongoResumePosition::ResumeAfter(token(encoded)?))
88        }
89        MongoCheckpointPosition::Stream(StreamPosition::StartAfter(encoded)) => {
90            ReaderStart::Stream(MongoResumePosition::StartAfter(token(encoded)?))
91        }
92        MongoCheckpointPosition::Stream(StreamPosition::StartAt(at)) => {
93            ReaderStart::Stream(MongoResumePosition::StartAt(*at))
94        }
95        MongoCheckpointPosition::Snapshot(cut) => ReaderStart::Snapshot(cut.clone()),
96    })
97}
98
99#[async_trait]
100impl SourceConnector for MongoDbCdcSource {
101    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
102        Some(self.task_tracker.clone())
103    }
104
105    fn recovery_identity_options(
106        &self,
107        config: &ConnectorConfig,
108    ) -> Result<Option<BTreeMap<String, String>>, ConnectorError> {
109        let mut parsed = parsed_config(&self.config, config)?;
110        parsed.normalize_pipeline()?;
111        let pipeline = super::super::config::canonical_pipeline_json(&parsed.pipeline);
112        let wire_protocol = match parsed.output_mode {
113            SourceOutputMode::History => "change-stream-history-v1",
114            SourceOutputMode::Document => "change-stream-document-v1",
115        };
116        Ok(Some(BTreeMap::from([
117            ("collection".into(), parsed.collection),
118            ("database".into(), parsed.database),
119            (
120                "full.document.mode".into(),
121                parsed.full_document_mode.to_string(),
122            ),
123            ("output.mode".into(), parsed.output_mode.to_string()),
124            ("snapshot.mode".into(), parsed.snapshot_mode.to_string()),
125            ("objectid.columns".into(), parsed.objectid_columns.join(",")),
126            (
127                "document.json.column".into(),
128                parsed.document_json_column.unwrap_or_default(),
129            ),
130            ("pipeline".into(), pipeline),
131            ("wire.protocol".into(), wire_protocol.into()),
132        ])))
133    }
134
135    async fn resolve_schema(
136        &mut self,
137        config: &ConnectorConfig,
138        explicit: Option<SchemaRef>,
139    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
140        super::schema_resolution::resolve(config, explicit).await
141    }
142
143    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
144        if self.state != ConnectorState::Created {
145            return Err(ConnectorError::InvalidState {
146                expected: ConnectorState::Created.to_string(),
147                actual: self.state.to_string(),
148            });
149        }
150        let (config, position, _) = request.into_parts();
151        let parsed = parsed_config(&self.config, &config)?;
152        let projection = document_projection(&parsed, &config, config.arrow_schema())?;
153        let (start, restored, mut expected_collection_uuid, mut expected_deployment_identity) =
154            match position {
155                SourcePosition::Initial => (
156                    match parsed.snapshot_mode {
157                        SnapshotMode::Never => ReaderStart::FreshChanges,
158                        SnapshotMode::Initial => ReaderStart::FreshSnapshot,
159                    },
160                    None,
161                    None,
162                    None,
163                ),
164                SourcePosition::Initialized { .. } => {
165                    return Err(ConnectorError::ConfigurationError(
166                        "MongoDB has no sealed topology startup contract".into(),
167                    ));
168                }
169                SourcePosition::Resume {
170                    attempt,
171                    checkpoint,
172                } => {
173                    let parsed_checkpoint = parse_mongodb_checkpoint(&checkpoint, &parsed)
174                        .map_err(|error| {
175                            ConnectorError::ConfigurationError(format!(
176                                "invalid MongoDB CDC checkpoint {attempt:?}: {error}"
177                            ))
178                        })?;
179                    (
180                        restored_start(&parsed_checkpoint)?,
181                        Some(parsed_checkpoint.emitted),
182                        Some(parsed_checkpoint.collection_uuid),
183                        Some(parsed_checkpoint.deployment_identity),
184                    )
185                }
186            };
187
188        if let Some((collection, deployment)) =
189            super::schema_resolution::committed_identity(config.schema_binding())?
190        {
191            if expected_collection_uuid.is_some_and(|expected| expected != collection)
192                || expected_deployment_identity
193                    .as_ref()
194                    .is_some_and(|expected| expected != &deployment)
195            {
196                return Err(ConnectorError::SchemaMismatch(
197                    "MongoDB checkpoint and committed collection identity disagree".into(),
198                ));
199            }
200            expected_collection_uuid = Some(collection);
201            expected_deployment_identity = Some(deployment);
202        }
203        self.start_change_stream_reader(ReaderLaunch {
204            config: parsed,
205            start,
206            restored,
207            expected_collection_uuid,
208            expected_deployment_identity,
209            projection,
210        })
211        .await?;
212
213        self.state = ConnectorState::Running;
214        tracing::info!(
215            database = %self.config.database,
216            collection = %self.config.collection,
217            output_mode = %self.config.output_mode,
218            snapshot_mode = %self.config.snapshot_mode,
219            full_document_mode = ?self.config.full_document_mode,
220            "MongoDB CDC source opened"
221        );
222
223        Ok(())
224    }
225
226    async fn poll_batch(
227        &mut self,
228        max_records: usize,
229    ) -> Result<Option<SourceBatch>, ConnectorError> {
230        self.drain_channel(max_records.saturating_sub(self.event_buffer.len()));
231        if let Some(batch) = self.drain_to_batch(max_records)? {
232            return Ok(Some(batch));
233        }
234        self.check_reader_error()?;
235        Ok(None)
236    }
237
238    fn schema(&self) -> SchemaRef {
239        Arc::clone(&self.schema)
240    }
241
242    fn checkpoint(&self) -> SourceCheckpoint {
243        let mut checkpoint = SourceCheckpoint::new();
244        // A configured namespace is not a physical replay identity until admission has read the
245        // server-assigned collection UUID and deployment identity.
246        let (Some(collection_uuid), Some(deployment_identity), Some(emitted)) = (
247            self.collection_uuid,
248            self.deployment_identity.as_ref(),
249            self.emitted.as_ref(),
250        ) else {
251            return checkpoint;
252        };
253        set_emitted_offsets(&mut checkpoint, emitted);
254        checkpoint.set_metadata("connector", MONGODB_CHECKPOINT_CONNECTOR);
255        checkpoint.set_metadata("version", MONGODB_CHECKPOINT_VERSION);
256        checkpoint.set_metadata("database", &self.config.database);
257        checkpoint.set_metadata("collection", &self.config.collection);
258        checkpoint.set_metadata(
259            COLLECTION_UUID_METADATA,
260            collection_uuid.hyphenated().to_string(),
261        );
262        checkpoint.set_metadata(DEPLOYMENT_IDENTITY_METADATA, deployment_identity.encode());
263        checkpoint.set_metadata(
264            STREAM_IDENTITY_METADATA,
265            mongodb_stream_identity(&self.config),
266        );
267        checkpoint
268    }
269
270    async fn notify_epoch_committed(
271        &mut self,
272        _epoch: u64,
273        checkpoint: &SourceCheckpoint,
274    ) -> Result<(), ConnectorError> {
275        let Some((at, _)) = self.snapshot_committed.as_ref() else {
276            return Ok(());
277        };
278        if checkpoint.offsets().is_empty() {
279            return Ok(());
280        }
281        let committed = parse_mongodb_checkpoint(checkpoint, &self.config)?;
282        if matches!(&committed.emitted.position, MongoCheckpointPosition::Snapshot(cut) if cut.at == *at)
283        {
284            if let Some((_, committed_tx)) = self.snapshot_committed.take() {
285                committed_tx.send_replace(true);
286            }
287        }
288        Ok(())
289    }
290
291    async fn close(&mut self) -> Result<(), ConnectorError> {
292        let mut reader_join_error = None;
293        if let Some(tx) = self.reader_shutdown.as_ref() {
294            tx.send_replace(true);
295        }
296        let mut detach_reader = false;
297        if let Some(handle) = self.reader_handle.as_mut() {
298            match tokio::time::timeout(READER_SHUTDOWN_TIMEOUT, &mut *handle).await {
299                Ok(Ok(())) => {}
300                Ok(Err(error)) if error.is_cancelled() => {}
301                Ok(Err(error)) => reader_join_error = Some(error.to_string()),
302                Err(_) => {
303                    tracing::warn!(
304                        "MongoDB CDC reader exceeded its close deadline; its tracked reaper retains shutdown ownership"
305                    );
306                    detach_reader = true;
307                }
308            }
309        }
310        if detach_reader {
311            if let Some(handle) = self.reader_handle.take() {
312                reap_mongo_reader(handle, &self.task_owner);
313            }
314        } else {
315            self.reader_handle = None;
316        }
317        self.reader_shutdown = None;
318        self.event_rx = None;
319        self.reader_error = None;
320        self.snapshot_committed = None;
321
322        self.event_buffer.clear();
323        self.state = ConnectorState::Closed;
324        tracing::info!("MongoDB CDC source closed");
325        if let Some(error) = reader_join_error {
326            return Err(ConnectorError::ReadError(format!(
327                "MongoDB CDC reader task failed during close: {error}"
328            )));
329        }
330        Ok(())
331    }
332
333    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
334        Some(Arc::clone(&self.data_ready))
335    }
336
337    fn contract(&self, config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
338        let parsed = parsed_config(&self.config, config)?;
339        // The initial copy starts only after its cut is durably committed.
340        let consistency = match parsed.snapshot_mode {
341            SnapshotMode::Never => SourceConsistency::Replayable,
342            SnapshotMode::Initial => SourceConsistency::CommitCoupled,
343        };
344        Ok(match parsed.output_mode {
345            SourceOutputMode::History => SourceContract::new(
346                consistency,
347                SourceTopology::Singleton,
348                SourceInputMode::AppendOnly,
349            ),
350            SourceOutputMode::Document => SourceContract::new(
351                consistency,
352                SourceTopology::Singleton,
353                SourceInputMode::KeyedUpsert,
354            )
355            .with_row_positions(SourceRowPositionCapability::OrderedDeterministic),
356        })
357    }
358}
laminar_connectors/mongodb/source/lifecycle.rs

laminar_connectors/mongodb/source/
lifecycle.rs
