laminar_connectors/mongodb/sink/
lifecycle.rs

1//! Sink admission, buffering lifecycle, flushing, and shutdown.
2
3use super::{
4    async_trait, info, Arc, ConnectorConfig, ConnectorError, ConnectorState, ConnectorTaskTracker,
5    Duration, MongoDbSink, RecordBatch, SchemaRef, SinkConnector, SinkConsistency, SinkContract,
6    SinkInputMode, SinkTopology, WriteMode, WriteResult, MAX_BUFFERED_RETAINED_BYTES,
7};
8
9#[async_trait]
10impl SinkConnector for MongoDbSink {
11    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
12        Some(self.task_tracker.clone())
13    }
14
15    fn contract(&self, config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
16        let (cfg, schema, _) = if config.properties().is_empty() {
17            (
18                self.config.clone(),
19                Arc::clone(&self.schema),
20                self.write_timeout,
21            )
22        } else {
23            Self::decode_connector_config(config)?
24        };
25        cfg.validate()?;
26        Self::validate_schema(&schema, &cfg)?;
27        let (topology, input_mode) = match cfg.write_mode {
28            WriteMode::Insert => (SinkTopology::MultiWriter, SinkInputMode::AppendOnly),
29            WriteMode::Upsert { .. } | WriteMode::CdcReplay => {
30                (SinkTopology::Singleton, SinkInputMode::FullChangelog)
31            }
32        };
33        Ok(SinkContract::new(
34            SinkConsistency::DurableAtLeastOnce,
35            topology,
36            input_mode,
37        ))
38    }
39
40    async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
41        if self.state != ConnectorState::Created {
42            return Err(ConnectorError::InvalidState {
43                expected: ConnectorState::Created.to_string(),
44                actual: self.state.to_string(),
45            });
46        }
47        if config.properties().is_empty() {
48            self.config.validate()?;
49            Self::validate_schema(&self.schema, &self.config)?;
50        } else {
51            self.apply_connector_config(config)?;
52        }
53        self.connect().await?;
54
55        self.state = ConnectorState::Running;
56        info!(
57            database = %self.config.database,
58            collection = %self.config.collection,
59            write_mode = ?self.config.write_mode,
60            "MongoDB sink opened"
61        );
62
63        Ok(())
64    }
65
66    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
67        self.write_batch_with_retained_limit(batch, MAX_BUFFERED_RETAINED_BYTES)
68            .await
69    }
70
71    fn schema(&self) -> SchemaRef {
72        Arc::clone(&self.schema)
73    }
74
75    fn suggested_write_timeout(&self) -> Duration {
76        self.write_timeout
77    }
78
79    fn flush_interval(&self) -> Duration {
80        self.config.flush_interval()
81    }
82
83    async fn flush(&mut self) -> Result<(), ConnectorError> {
84        if self.state != ConnectorState::Running {
85            return Err(ConnectorError::InvalidState {
86                expected: ConnectorState::Running.to_string(),
87                actual: self.state.to_string(),
88            });
89        }
90        self.flush_inner().await.map(|_| ())
91    }
92
93    async fn close(&mut self) -> Result<(), ConnectorError> {
94        if self.state == ConnectorState::Closed {
95            return Ok(());
96        }
97
98        let flush_result = if self.state == ConnectorState::Running && !self.buffer.is_empty() {
99            self.flush_inner().await.map(|_| ())
100        } else {
101            drop(self.take_buffer());
102            Ok(())
103        };
104        self.collection = None;
105        self.client = None;
106        self.state = ConnectorState::Closed;
107        info!("MongoDB sink closed");
108        flush_result
109    }
110}
laminar_connectors/mongodb/sink/lifecycle.rs

laminar_connectors/mongodb/sink/
lifecycle.rs
