risingwave_connector/source/
mod.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod prelude {
16    // import all split enumerators
17    pub use crate::source::datagen::DatagenSplitEnumerator;
18    pub use crate::source::filesystem::LegacyS3SplitEnumerator;
19    pub use crate::source::filesystem::opendal_source::OpendalEnumerator;
20    pub use crate::source::google_pubsub::PubsubSplitEnumerator as GooglePubsubSplitEnumerator;
21    pub use crate::source::iceberg::IcebergSplitEnumerator;
22    pub use crate::source::kafka::KafkaSplitEnumerator;
23    pub use crate::source::kinesis::KinesisSplitEnumerator;
24    pub use crate::source::mqtt::MqttSplitEnumerator;
25    pub use crate::source::nats::NatsSplitEnumerator;
26    pub use crate::source::nexmark::NexmarkSplitEnumerator;
27    pub use crate::source::pulsar::PulsarSplitEnumerator;
28    pub use crate::source::test_source::TestSourceSplitEnumerator as TestSplitEnumerator;
29    pub type AzblobSplitEnumerator =
30        OpendalEnumerator<crate::source::filesystem::opendal_source::OpendalAzblob>;
31    pub type GcsSplitEnumerator =
32        OpendalEnumerator<crate::source::filesystem::opendal_source::OpendalGcs>;
33    pub type OpendalS3SplitEnumerator =
34        OpendalEnumerator<crate::source::filesystem::opendal_source::OpendalS3>;
35    pub type PosixFsSplitEnumerator =
36        OpendalEnumerator<crate::source::filesystem::opendal_source::OpendalPosixFs>;
37    pub use crate::source::cdc::enumerator::DebeziumSplitEnumerator;
38    pub use crate::source::filesystem::opendal_source::BatchPosixFsEnumerator as BatchPosixFsSplitEnumerator;
39    pub type CitusCdcSplitEnumerator = DebeziumSplitEnumerator<crate::source::cdc::Citus>;
40    pub type MongodbCdcSplitEnumerator = DebeziumSplitEnumerator<crate::source::cdc::Mongodb>;
41    pub type PostgresCdcSplitEnumerator = DebeziumSplitEnumerator<crate::source::cdc::Postgres>;
42    pub type MysqlCdcSplitEnumerator = DebeziumSplitEnumerator<crate::source::cdc::Mysql>;
43    pub type SqlServerCdcSplitEnumerator = DebeziumSplitEnumerator<crate::source::cdc::SqlServer>;
44}
45
46pub mod base;
47pub mod batch;
48pub mod cdc;
49pub mod data_gen_util;
50pub mod datagen;
51pub mod filesystem;
52pub mod google_pubsub;
53pub mod kafka;
54pub mod kinesis;
55pub mod monitor;
56pub mod mqtt;
57pub mod nats;
58pub mod nexmark;
59pub mod pulsar;
60pub mod utils;
61
62mod util;
63use std::collections::HashMap;
64use std::future::IntoFuture;
65use std::time::Duration;
66
67pub use base::{UPSTREAM_SOURCE_KEY, WEBHOOK_CONNECTOR, *};
68pub use batch::BatchSourceSplitImpl;
69pub(crate) use common::*;
70use google_cloud_pubsub::subscription::Subscription;
71pub use google_pubsub::GOOGLE_PUBSUB_CONNECTOR;
72pub use kafka::KAFKA_CONNECTOR;
73pub use kinesis::KINESIS_CONNECTOR;
74use monitor::{ConnectorAckFailureType, GLOBAL_SOURCE_METRICS};
75pub use mqtt::MQTT_CONNECTOR;
76pub use nats::NATS_CONNECTOR;
77use utils::feature_gated_source_mod;
78
79pub use self::adbc_snowflake::ADBC_SNOWFLAKE_CONNECTOR;
80mod common;
81pub mod iceberg;
82mod manager;
83pub mod reader;
84pub mod test_source;
85feature_gated_source_mod!(adbc_snowflake, "adbc_snowflake");
86
87use async_nats::jetstream::consumer::AckPolicy as JetStreamAckPolicy;
88use async_nats::jetstream::context::Context as JetStreamContext;
89pub use manager::{SourceColumnDesc, SourceColumnType};
90use risingwave_common::array::{Array, ArrayRef};
91use risingwave_common::row::OwnedRow;
92use risingwave_pb::id::{ActorId, SourceId};
93use thiserror_ext::AsReport;
94pub use util::fill_adaptive_split;
95
96pub use crate::source::filesystem::LEGACY_S3_CONNECTOR;
97pub use crate::source::filesystem::opendal_source::{
98    AZBLOB_CONNECTOR, BATCH_POSIX_FS_CONNECTOR, GCS_CONNECTOR, OPENDAL_S3_CONNECTOR,
99    POSIX_FS_CONNECTOR,
100};
101pub use crate::source::nexmark::NEXMARK_CONNECTOR;
102pub use crate::source::pulsar::PULSAR_CONNECTOR;
103use crate::source::pulsar::source::reader::PULSAR_ACK_CHANNEL;
104
105pub fn should_copy_to_format_encode_options(key: &str, connector: &str) -> bool {
106    const PREFIXES: &[&str] = &[
107        "schema.registry",
108        "schema.location",
109        "message",
110        "key.message",
111        "without_header",
112        "delimiter",
113        // AwsAuthProps
114        "region",
115        "endpoint_url",
116        "access_key",
117        "secret_key",
118        "session_token",
119        "arn",
120        "external_id",
121        "profile",
122    ];
123    PREFIXES.iter().any(|prefix| key.starts_with(prefix))
124        || (key == "endpoint" && !connector.eq_ignore_ascii_case(KINESIS_CONNECTOR))
125}
126
127/// Tasks executed by `WaitCheckpointWorker`
128pub enum WaitCheckpointTask {
129    CommitCdcOffset(Option<(SplitId, String)>),
130    AckPubsubMessage(Subscription, Vec<ArrayRef>),
131    AckNatsJetStream(JetStreamContext, Vec<ArrayRef>, JetStreamAckPolicy),
132    AckPulsarMessage(Vec<(String, ArrayRef)>),
133}
134
135impl WaitCheckpointTask {
136    /// Create a fresh task for the next epoch, reusing expensive-to-create clients
137    /// (e.g. `PubSub` `Subscription`, NATS `JetStreamContext`) from the current task.
138    /// This avoids re-establishing gRPC/network connections on every checkpoint.
139    pub fn reset_for_next_epoch(&self) -> Self {
140        match self {
141            WaitCheckpointTask::CommitCdcOffset(_) => WaitCheckpointTask::CommitCdcOffset(None),
142            WaitCheckpointTask::AckPubsubMessage(subscription, _) => {
143                WaitCheckpointTask::AckPubsubMessage(subscription.clone(), vec![])
144            }
145            WaitCheckpointTask::AckNatsJetStream(context, _, ack_policy) => {
146                WaitCheckpointTask::AckNatsJetStream(context.clone(), vec![], *ack_policy)
147            }
148            WaitCheckpointTask::AckPulsarMessage(_) => WaitCheckpointTask::AckPulsarMessage(vec![]),
149        }
150    }
151
152    pub async fn run(self, source_id: SourceId, source_name: &str) {
153        self.run_with_on_commit_success(source_id, source_name, |_source_id, _offset| {
154            // Default implementation: no action on commit success
155        })
156        .await;
157    }
158
159    pub async fn run_with_on_commit_success<F>(
160        self,
161        source_id: SourceId,
162        source_name: &str,
163        mut on_commit_success: F,
164    ) where
165        F: FnMut(u64, &str),
166    {
167        use std::str::FromStr;
168        let source_id_label = source_id.to_string();
169        match self {
170            WaitCheckpointTask::CommitCdcOffset(updated_offset) => {
171                if let Some((split_id, offset)) = updated_offset {
172                    let committed_source_id: u64 = u64::from_str(split_id.as_ref()).unwrap();
173                    // notify cdc connector to commit offset
174                    match cdc::jni_source::commit_cdc_offset(committed_source_id, offset.clone()) {
175                        Ok(()) => {
176                            // Execute callback after successful commit
177                            on_commit_success(committed_source_id, &offset);
178                        }
179                        Err(e) => {
180                            tracing::error!(
181                                source_id = committed_source_id,
182                                source_name,
183                                error = %e.as_report(),
184                                "source#{committed_source_id}: failed to commit cdc offset: {offset}.",
185                            )
186                        }
187                    }
188                }
189            }
190            WaitCheckpointTask::AckPulsarMessage(ack_array) => {
191                let mut latest_ack_by_channel = HashMap::new();
192                for (ack_channel_id, to_cumulative_ack) in ack_array {
193                    let encode_message_id_data = to_cumulative_ack
194                        .as_bytea()
195                        .iter()
196                        .flatten()
197                        .last()
198                        .map(|message_id| message_id.to_owned());
199
200                    if let Some(encode_message_id_data) = encode_message_id_data {
201                        latest_ack_by_channel.insert(ack_channel_id, Some(encode_message_id_data));
202                    } else {
203                        latest_ack_by_channel.entry(ack_channel_id).or_insert(None);
204                    }
205                }
206
207                for (ack_channel_id, encode_message_id_data) in latest_ack_by_channel {
208                    let Some(encode_message_id_data) = encode_message_id_data else {
209                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
210                            source_name,
211                            "pulsar",
212                            ConnectorAckFailureType::EmptyMessageId,
213                        );
214                        tracing::warn!(
215                            source_id = source_id_label,
216                            source_name,
217                            ack_channel_id,
218                            "skip Pulsar ack because the checkpoint ack batches have no message id",
219                        );
220                        continue;
221                    };
222
223                    let Some(ack_tx) = PULSAR_ACK_CHANNEL.lock().get(&ack_channel_id).cloned()
224                    else {
225                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
226                            source_name,
227                            "pulsar",
228                            ConnectorAckFailureType::ChannelMissing,
229                        );
230                        tracing::warn!(
231                            source_id = source_id_label,
232                            source_name,
233                            ack_channel_id,
234                            "skip Pulsar ack because the ack channel is missing",
235                        );
236                        continue;
237                    };
238
239                    if let Err(e) = ack_tx.send(encode_message_id_data) {
240                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
241                            source_name,
242                            "pulsar",
243                            ConnectorAckFailureType::ChannelSendError,
244                        );
245                        tracing::warn!(
246                            source_id = source_id_label,
247                            source_name,
248                            ack_channel_id,
249                            error = %e.as_report(),
250                            "failed to send Pulsar ack message id to the reader ack channel",
251                        );
252                    }
253                }
254            }
255            WaitCheckpointTask::AckPubsubMessage(subscription, ack_id_arrs) => {
256                const ACK_RPC_TIMEOUT: Duration = Duration::from_secs(30);
257                async fn ack(
258                    subscription: &Subscription,
259                    ack_ids: Vec<String>,
260                    source_id_label: &str,
261                    source_name: &str,
262                ) {
263                    if ack_ids.is_empty() {
264                        return;
265                    }
266                    tracing::trace!("acking pubsub messages {:?}", ack_ids);
267                    match tokio::time::timeout(ACK_RPC_TIMEOUT, subscription.ack(ack_ids)).await {
268                        Ok(Ok(())) => {
269                            GLOBAL_SOURCE_METRICS
270                                .inc_connector_ack_success_count(source_name, "pubsub");
271                        }
272                        Ok(Err(e)) => {
273                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
274                                source_name,
275                                "pubsub",
276                                ConnectorAckFailureType::Error,
277                            );
278                            tracing::error!(
279                                source_id = source_id_label,
280                                source_name,
281                                error = %e.as_report(),
282                                "failed to ack pubsub messages",
283                            )
284                        }
285                        Err(_) => {
286                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
287                                source_name,
288                                "pubsub",
289                                ConnectorAckFailureType::Timeout,
290                            );
291                            tracing::error!(
292                                source_id = source_id_label,
293                                source_name,
294                                "pubsub ack timed out after {ACK_RPC_TIMEOUT:?}",
295                            )
296                        }
297                    }
298                }
299                const MAX_ACK_BATCH_SIZE: usize = 1000;
300                let mut ack_ids: Vec<String> = vec![];
301                for arr in ack_id_arrs {
302                    for ack_id in arr.as_utf8().iter().flatten() {
303                        ack_ids.push(ack_id.to_owned());
304                        if ack_ids.len() >= MAX_ACK_BATCH_SIZE {
305                            ack(
306                                &subscription,
307                                std::mem::take(&mut ack_ids),
308                                &source_id_label,
309                                source_name,
310                            )
311                            .await;
312                        }
313                    }
314                }
315                ack(&subscription, ack_ids, &source_id_label, source_name).await;
316            }
317            WaitCheckpointTask::AckNatsJetStream(
318                ref context,
319                reply_subjects_arrs,
320                ref ack_policy,
321            ) => {
322                const ACK_RPC_TIMEOUT: Duration = Duration::from_secs(30);
323                async fn ack(
324                    context: &JetStreamContext,
325                    reply_subject: String,
326                    source_id_label: &str,
327                    source_name: &str,
328                ) {
329                    let fut = async {
330                        let ack_future = context
331                            .publish(reply_subject.clone(), "+ACK".into())
332                            .await
333                            .map_err(|e| e.to_report_string())?;
334                        ack_future
335                            .into_future()
336                            .await
337                            .map_err(|e| e.to_report_string())?;
338                        Ok::<(), String>(())
339                    };
340                    match tokio::time::timeout(ACK_RPC_TIMEOUT, fut).await {
341                        Ok(Ok(())) => {
342                            GLOBAL_SOURCE_METRICS
343                                .inc_connector_ack_success_count(source_name, "nats_jetstream");
344                        }
345                        Ok(Err(e)) => {
346                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
347                                source_name,
348                                "nats_jetstream",
349                                ConnectorAckFailureType::Error,
350                            );
351                            tracing::error!(
352                                source_id = source_id_label,
353                                source_name,
354                                error = %e,
355                                subject = ?reply_subject,
356                                "failed to ack NATS JetStream message",
357                            );
358                        }
359                        Err(_) => {
360                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
361                                source_name,
362                                "nats_jetstream",
363                                ConnectorAckFailureType::Timeout,
364                            );
365                            tracing::error!(
366                                source_id = source_id_label,
367                                source_name,
368                                subject = ?reply_subject,
369                                "NATS JetStream ack timed out after {ACK_RPC_TIMEOUT:?}",
370                            );
371                        }
372                    }
373                }
374
375                let reply_subjects = reply_subjects_arrs
376                    .iter()
377                    .flat_map(|arr| {
378                        arr.as_utf8()
379                            .iter()
380                            .flatten()
381                            .map(|s| s.to_owned())
382                            .collect::<Vec<String>>()
383                    })
384                    .collect::<Vec<String>>();
385
386                match ack_policy {
387                    JetStreamAckPolicy::None | JetStreamAckPolicy::FlowControl => (),
388                    JetStreamAckPolicy::Explicit => {
389                        for reply_subject in reply_subjects {
390                            if reply_subject.is_empty() {
391                                continue;
392                            }
393                            ack(context, reply_subject, &source_id_label, source_name).await;
394                        }
395                    }
396                    JetStreamAckPolicy::All => {
397                        if let Some(reply_subject) = reply_subjects.last() {
398                            ack(
399                                context,
400                                reply_subject.clone(),
401                                &source_id_label,
402                                source_name,
403                            )
404                            .await;
405                        }
406                    }
407                }
408            }
409        }
410    }
411}
412
413#[derive(Clone, Debug, PartialEq)]
414pub struct CdcTableSnapshotSplitCommon<T: Clone> {
415    pub split_id: i64,
416    pub left_bound_inclusive: T,
417    pub right_bound_exclusive: T,
418}
419
420pub type CdcTableSnapshotSplit = CdcTableSnapshotSplitCommon<OwnedRow>;
421pub type CdcTableSnapshotSplitRaw = CdcTableSnapshotSplitCommon<Vec<u8>>;
422
423/// Build the identifier of the ACK channel for a Pulsar reader.
424///
425/// Multiple actors can consume the same source split in one process, so the actor ID is required
426/// to prevent one reader from replacing another reader's channel.
427#[inline]
428pub fn build_pulsar_ack_channel_id(
429    source_id: SourceId,
430    split_id: &SplitId,
431    actor_id: ActorId,
432) -> String {
433    format!("{}-{}-{}", source_id, split_id, actor_id)
434}
435
436#[cfg(test)]
437mod tests {
438    use risingwave_common::array::{Array, BytesArray};
439    use tokio::sync::mpsc::error::TryRecvError;
440
441    use super::*;
442
443    #[test]
444    fn test_pulsar_ack_channel_id_is_actor_scoped() {
445        let source_id = SourceId::new(7);
446        let split_id: SplitId = "persistent://public/default/topic".into();
447
448        let first = build_pulsar_ack_channel_id(source_id, &split_id, ActorId::new(11));
449        let second = build_pulsar_ack_channel_id(source_id, &split_id, ActorId::new(12));
450
451        assert_ne!(first, second);
452    }
453
454    fn message_ids<const N: usize>(values: [Option<&[u8]>; N]) -> ArrayRef {
455        BytesArray::from_iter(values).into_ref()
456    }
457
458    #[tokio::test]
459    async fn test_ack_pulsar_message_for_each_split() {
460        let split_0_channel = "test-pulsar-ack-multiple-splits-0".to_owned();
461        let split_1_channel = "test-pulsar-ack-multiple-splits-1".to_owned();
462        let empty_split_channel = "test-pulsar-ack-multiple-splits-empty".to_owned();
463        let (split_0_tx, mut split_0_rx) = tokio::sync::mpsc::unbounded_channel();
464        let (split_1_tx, mut split_1_rx) = tokio::sync::mpsc::unbounded_channel();
465        let (empty_split_tx, mut empty_split_rx) = tokio::sync::mpsc::unbounded_channel();
466
467        PULSAR_ACK_CHANNEL
468            .lock()
469            .insert(split_0_channel.clone(), split_0_tx);
470        PULSAR_ACK_CHANNEL
471            .lock()
472            .insert(split_1_channel.clone(), split_1_tx);
473        PULSAR_ACK_CHANNEL
474            .lock()
475            .insert(empty_split_channel.clone(), empty_split_tx);
476
477        WaitCheckpointTask::AckPulsarMessage(vec![
478            (split_0_channel.clone(), message_ids([Some(b"split-0-old")])),
479            (
480                split_1_channel.clone(),
481                message_ids([Some(b"split-1-latest"), None]),
482            ),
483            (
484                split_0_channel.clone(),
485                message_ids([Some(b"split-0-latest")]),
486            ),
487            (split_0_channel.clone(), message_ids([None])),
488            (empty_split_channel.clone(), message_ids([None, None])),
489        ])
490        .run(SourceId::new(26891), "test_pulsar_source")
491        .await;
492
493        assert_eq!(split_0_rx.try_recv().unwrap(), b"split-0-latest");
494        assert_eq!(split_0_rx.try_recv(), Err(TryRecvError::Empty));
495        assert_eq!(split_1_rx.try_recv().unwrap(), b"split-1-latest");
496        assert_eq!(split_1_rx.try_recv(), Err(TryRecvError::Empty));
497        assert_eq!(empty_split_rx.try_recv(), Err(TryRecvError::Empty));
498
499        let mut ack_channels = PULSAR_ACK_CHANNEL.lock();
500        ack_channels.remove(&split_0_channel);
501        ack_channels.remove(&split_1_channel);
502        ack_channels.remove(&empty_split_channel);
503    }
504}
risingwave_connector/source/mod.rs

risingwave_connector/source/
mod.rs
