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

risingwave_connector/source/
mod.rs
