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        // `cdc_replay` interprets immutable history records itself; it consumes no engine
28        // changelog, so it advertises append-only input and admission rejects retractions.
29        let (topology, input_mode) = match cfg.write_mode {
30            WriteMode::Insert => (SinkTopology::MultiWriter, SinkInputMode::AppendOnly),
31            WriteMode::CdcReplay => (SinkTopology::Singleton, SinkInputMode::AppendOnly),
32            WriteMode::Upsert { .. } => (SinkTopology::Singleton, SinkInputMode::FullChangelog),
33        };
34        Ok(SinkContract::new(
35            SinkConsistency::DurableAtLeastOnce,
36            topology,
37            input_mode,
38        ))
39    }
40
41    async fn resolve_schema(
42        &mut self,
43        config: &ConnectorConfig,
44        input: SchemaRef,
45    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
46        let parsed = super::super::config::MongoDbSinkConfig::from_config(config)?;
47        super::schema_resolution::resolve(config, &parsed, input).await
48    }
49
50    async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
51        if self.state != ConnectorState::Created {
52            return Err(ConnectorError::InvalidState {
53                expected: ConnectorState::Created.to_string(),
54                actual: self.state.to_string(),
55            });
56        }
57        if config.properties().is_empty() {
58            self.config.validate()?;
59            Self::validate_schema(&self.schema, &self.config)?;
60        } else {
61            self.apply_connector_config(config)?;
62        }
63        let mut current =
64            super::schema_resolution::resolve(config, &self.config, self.schema.clone()).await?;
65        if config
66            .schema_binding()
67            .is_some_and(|binding| binding.value != current.value)
68        {
69            return Err(ConnectorError::SchemaMismatch("MongoDB collection identity or validator changed before activation; migrate the binding".into()));
70        }
71        if config.schema_binding().is_none() {
72            self.prepare_target(config, &mut current).await?;
73        }
74        if current.value.is_none() {
75            return Err(ConnectorError::SchemaMismatch(
76                "MongoDB activation requires a prepared collection UUID".into(),
77            ));
78        }
79        self.binding = Some(current);
80        self.connect().await?;
81
82        self.state = ConnectorState::Running;
83        info!(
84            database = %self.config.database,
85            collection = %self.config.collection,
86            write_mode = ?self.config.write_mode,
87            "MongoDB sink opened"
88        );
89
90        Ok(())
91    }
92
93    async fn prepare_schema(
94        &mut self,
95        config: &ConnectorConfig,
96        binding: &mut crate::schema::resolution::SchemaBinding,
97    ) -> Result<(), ConnectorError> {
98        self.prepare_target(config, binding).await
99    }
100
101    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
102        self.write_batch_with_retained_limit(batch, MAX_BUFFERED_RETAINED_BYTES)
103            .await
104    }
105
106    fn schema(&self) -> SchemaRef {
107        Arc::clone(&self.schema)
108    }
109
110    fn suggested_write_timeout(&self) -> Duration {
111        self.write_timeout
112    }
113
114    fn flush_interval(&self) -> Duration {
115        self.config.flush_interval()
116    }
117
118    async fn flush(&mut self) -> Result<(), ConnectorError> {
119        if self.state != ConnectorState::Running {
120            return Err(ConnectorError::InvalidState {
121                expected: ConnectorState::Running.to_string(),
122                actual: self.state.to_string(),
123            });
124        }
125        self.flush_inner().await.map(|_| ())
126    }
127
128    async fn close(&mut self) -> Result<(), ConnectorError> {
129        if self.state == ConnectorState::Closed {
130            return Ok(());
131        }
132
133        let flush_result = if self.state == ConnectorState::Running && !self.buffer.is_empty() {
134            self.flush_inner().await.map(|_| ())
135        } else {
136            drop(self.take_buffer());
137            Ok(())
138        };
139        self.collection = None;
140        self.client = None;
141        self.state = ConnectorState::Closed;
142        info!("MongoDB sink closed");
143        flush_result
144    }
145}
laminar_connectors/mongodb/sink/lifecycle.rs

laminar_connectors/mongodb/sink/
lifecycle.rs
