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::future::IntoFuture;
64use std::time::Duration;
65
66pub use base::{UPSTREAM_SOURCE_KEY, WEBHOOK_CONNECTOR, *};
67pub use batch::BatchSourceSplitImpl;
68pub(crate) use common::*;
69use google_cloud_pubsub::subscription::Subscription;
70pub use google_pubsub::GOOGLE_PUBSUB_CONNECTOR;
71pub use kafka::KAFKA_CONNECTOR;
72pub use kinesis::KINESIS_CONNECTOR;
73use monitor::{ConnectorAckFailureType, GLOBAL_SOURCE_METRICS};
74pub use mqtt::MQTT_CONNECTOR;
75pub use nats::NATS_CONNECTOR;
76use utils::feature_gated_source_mod;
77
78pub use self::adbc_snowflake::ADBC_SNOWFLAKE_CONNECTOR;
79mod common;
80pub mod iceberg;
81mod manager;
82pub mod reader;
83pub mod test_source;
84feature_gated_source_mod!(adbc_snowflake, "adbc_snowflake");
85
86use async_nats::jetstream::consumer::AckPolicy as JetStreamAckPolicy;
87use async_nats::jetstream::context::Context as JetStreamContext;
88pub use manager::{SourceColumnDesc, SourceColumnType};
89use risingwave_common::array::{Array, ArrayRef};
90use risingwave_common::row::OwnedRow;
91use risingwave_pb::id::SourceId;
92use thiserror_ext::AsReport;
93pub use util::fill_adaptive_split;
94
95pub use crate::source::filesystem::LEGACY_S3_CONNECTOR;
96pub use crate::source::filesystem::opendal_source::{
97    AZBLOB_CONNECTOR, BATCH_POSIX_FS_CONNECTOR, GCS_CONNECTOR, OPENDAL_S3_CONNECTOR,
98    POSIX_FS_CONNECTOR,
99};
100pub use crate::source::nexmark::NEXMARK_CONNECTOR;
101pub use crate::source::pulsar::PULSAR_CONNECTOR;
102use crate::source::pulsar::source::reader::PULSAR_ACK_CHANNEL;
103
104pub fn should_copy_to_format_encode_options(key: &str, connector: &str) -> bool {
105    const PREFIXES: &[&str] = &[
106        "schema.registry",
107        "schema.location",
108        "message",
109        "key.message",
110        "without_header",
111        "delimiter",
112        // AwsAuthProps
113        "region",
114        "endpoint_url",
115        "access_key",
116        "secret_key",
117        "session_token",
118        "arn",
119        "external_id",
120        "profile",
121    ];
122    PREFIXES.iter().any(|prefix| key.starts_with(prefix))
123        || (key == "endpoint" && !connector.eq_ignore_ascii_case(KINESIS_CONNECTOR))
124}
125
126/// Tasks executed by `WaitCheckpointWorker`
127pub enum WaitCheckpointTask {
128    CommitCdcOffset(Option<(SplitId, String)>),
129    AckPubsubMessage(Subscription, Vec<ArrayRef>),
130    AckNatsJetStream(JetStreamContext, Vec<ArrayRef>, JetStreamAckPolicy),
131    AckPulsarMessage(Vec<(String, ArrayRef)>),
132}
133
134impl WaitCheckpointTask {
135    /// Create a fresh task for the next epoch, reusing expensive-to-create clients
136    /// (e.g. `PubSub` `Subscription`, NATS `JetStreamContext`) from the current task.
137    /// This avoids re-establishing gRPC/network connections on every checkpoint.
138    pub fn reset_for_next_epoch(&self) -> Self {
139        match self {
140            WaitCheckpointTask::CommitCdcOffset(_) => WaitCheckpointTask::CommitCdcOffset(None),
141            WaitCheckpointTask::AckPubsubMessage(subscription, _) => {
142                WaitCheckpointTask::AckPubsubMessage(subscription.clone(), vec![])
143            }
144            WaitCheckpointTask::AckNatsJetStream(context, _, ack_policy) => {
145                WaitCheckpointTask::AckNatsJetStream(context.clone(), vec![], *ack_policy)
146            }
147            WaitCheckpointTask::AckPulsarMessage(_) => WaitCheckpointTask::AckPulsarMessage(vec![]),
148        }
149    }
150
151    pub async fn run(self, source_id: SourceId, source_name: &str) {
152        self.run_with_on_commit_success(source_id, source_name, |_source_id, _offset| {
153            // Default implementation: no action on commit success
154        })
155        .await;
156    }
157
158    pub async fn run_with_on_commit_success<F>(
159        self,
160        source_id: SourceId,
161        source_name: &str,
162        mut on_commit_success: F,
163    ) where
164        F: FnMut(u64, &str),
165    {
166        use std::str::FromStr;
167        let source_id_label = source_id.to_string();
168        match self {
169            WaitCheckpointTask::CommitCdcOffset(updated_offset) => {
170                if let Some((split_id, offset)) = updated_offset {
171                    let committed_source_id: u64 = u64::from_str(split_id.as_ref()).unwrap();
172                    // notify cdc connector to commit offset
173                    match cdc::jni_source::commit_cdc_offset(committed_source_id, offset.clone()) {
174                        Ok(()) => {
175                            // Execute callback after successful commit
176                            on_commit_success(committed_source_id, &offset);
177                        }
178                        Err(e) => {
179                            tracing::error!(
180                                source_id = committed_source_id,
181                                source_name,
182                                error = %e.as_report(),
183                                "source#{committed_source_id}: failed to commit cdc offset: {offset}.",
184                            )
185                        }
186                    }
187                }
188            }
189            WaitCheckpointTask::AckPulsarMessage(ack_array) => {
190                if let Some((ack_channel_id, to_cumulative_ack)) = ack_array.last() {
191                    let Some(encode_message_id_data) = to_cumulative_ack
192                        .as_bytea()
193                        .iter()
194                        .last()
195                        .flatten()
196                        .map(|x| x.to_owned())
197                    else {
198                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
199                            source_name,
200                            "pulsar",
201                            ConnectorAckFailureType::EmptyMessageId,
202                        );
203                        tracing::warn!(
204                            source_id = source_id_label,
205                            source_name,
206                            ack_channel_id,
207                            "skip Pulsar ack because the checkpoint ack batch has no message id",
208                        );
209                        return;
210                    };
211
212                    let Some(ack_tx) = PULSAR_ACK_CHANNEL.get(ack_channel_id).await else {
213                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
214                            source_name,
215                            "pulsar",
216                            ConnectorAckFailureType::ChannelMissing,
217                        );
218                        tracing::warn!(
219                            source_id = source_id_label,
220                            source_name,
221                            ack_channel_id,
222                            "skip Pulsar ack because the ack channel is missing",
223                        );
224                        return;
225                    };
226
227                    if let Err(e) = ack_tx.send(encode_message_id_data) {
228                        GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
229                            source_name,
230                            "pulsar",
231                            ConnectorAckFailureType::ChannelSendError,
232                        );
233                        tracing::warn!(
234                            source_id = source_id_label,
235                            source_name,
236                            ack_channel_id,
237                            error = %e.as_report(),
238                            "failed to send Pulsar ack message id to the reader ack channel",
239                        );
240                    }
241                }
242            }
243            WaitCheckpointTask::AckPubsubMessage(subscription, ack_id_arrs) => {
244                const ACK_RPC_TIMEOUT: Duration = Duration::from_secs(30);
245                async fn ack(
246                    subscription: &Subscription,
247                    ack_ids: Vec<String>,
248                    source_id_label: &str,
249                    source_name: &str,
250                ) {
251                    if ack_ids.is_empty() {
252                        return;
253                    }
254                    tracing::trace!("acking pubsub messages {:?}", ack_ids);
255                    match tokio::time::timeout(ACK_RPC_TIMEOUT, subscription.ack(ack_ids)).await {
256                        Ok(Ok(())) => {
257                            GLOBAL_SOURCE_METRICS
258                                .inc_connector_ack_success_count(source_name, "pubsub");
259                        }
260                        Ok(Err(e)) => {
261                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
262                                source_name,
263                                "pubsub",
264                                ConnectorAckFailureType::Error,
265                            );
266                            tracing::error!(
267                                source_id = source_id_label,
268                                source_name,
269                                error = %e.as_report(),
270                                "failed to ack pubsub messages",
271                            )
272                        }
273                        Err(_) => {
274                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
275                                source_name,
276                                "pubsub",
277                                ConnectorAckFailureType::Timeout,
278                            );
279                            tracing::error!(
280                                source_id = source_id_label,
281                                source_name,
282                                "pubsub ack timed out after {ACK_RPC_TIMEOUT:?}",
283                            )
284                        }
285                    }
286                }
287                const MAX_ACK_BATCH_SIZE: usize = 1000;
288                let mut ack_ids: Vec<String> = vec![];
289                for arr in ack_id_arrs {
290                    for ack_id in arr.as_utf8().iter().flatten() {
291                        ack_ids.push(ack_id.to_owned());
292                        if ack_ids.len() >= MAX_ACK_BATCH_SIZE {
293                            ack(
294                                &subscription,
295                                std::mem::take(&mut ack_ids),
296                                &source_id_label,
297                                source_name,
298                            )
299                            .await;
300                        }
301                    }
302                }
303                ack(&subscription, ack_ids, &source_id_label, source_name).await;
304            }
305            WaitCheckpointTask::AckNatsJetStream(
306                ref context,
307                reply_subjects_arrs,
308                ref ack_policy,
309            ) => {
310                const ACK_RPC_TIMEOUT: Duration = Duration::from_secs(30);
311                async fn ack(
312                    context: &JetStreamContext,
313                    reply_subject: String,
314                    source_id_label: &str,
315                    source_name: &str,
316                ) {
317                    let fut = async {
318                        let ack_future = context
319                            .publish(reply_subject.clone(), "+ACK".into())
320                            .await
321                            .map_err(|e| e.to_report_string())?;
322                        ack_future
323                            .into_future()
324                            .await
325                            .map_err(|e| e.to_report_string())?;
326                        Ok::<(), String>(())
327                    };
328                    match tokio::time::timeout(ACK_RPC_TIMEOUT, fut).await {
329                        Ok(Ok(())) => {
330                            GLOBAL_SOURCE_METRICS
331                                .inc_connector_ack_success_count(source_name, "nats_jetstream");
332                        }
333                        Ok(Err(e)) => {
334                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
335                                source_name,
336                                "nats_jetstream",
337                                ConnectorAckFailureType::Error,
338                            );
339                            tracing::error!(
340                                source_id = source_id_label,
341                                source_name,
342                                error = %e,
343                                subject = ?reply_subject,
344                                "failed to ack NATS JetStream message",
345                            );
346                        }
347                        Err(_) => {
348                            GLOBAL_SOURCE_METRICS.inc_connector_ack_failure_count(
349                                source_name,
350                                "nats_jetstream",
351                                ConnectorAckFailureType::Timeout,
352                            );
353                            tracing::error!(
354                                source_id = source_id_label,
355                                source_name,
356                                subject = ?reply_subject,
357                                "NATS JetStream ack timed out after {ACK_RPC_TIMEOUT:?}",
358                            );
359                        }
360                    }
361                }
362
363                let reply_subjects = reply_subjects_arrs
364                    .iter()
365                    .flat_map(|arr| {
366                        arr.as_utf8()
367                            .iter()
368                            .flatten()
369                            .map(|s| s.to_owned())
370                            .collect::<Vec<String>>()
371                    })
372                    .collect::<Vec<String>>();
373
374                match ack_policy {
375                    JetStreamAckPolicy::None => (),
376                    JetStreamAckPolicy::Explicit => {
377                        for reply_subject in reply_subjects {
378                            if reply_subject.is_empty() {
379                                continue;
380                            }
381                            ack(context, reply_subject, &source_id_label, source_name).await;
382                        }
383                    }
384                    JetStreamAckPolicy::All => {
385                        if let Some(reply_subject) = reply_subjects.last() {
386                            ack(
387                                context,
388                                reply_subject.clone(),
389                                &source_id_label,
390                                source_name,
391                            )
392                            .await;
393                        }
394                    }
395                }
396            }
397        }
398    }
399}
400
401#[derive(Clone, Debug, PartialEq)]
402pub struct CdcTableSnapshotSplitCommon<T: Clone> {
403    pub split_id: i64,
404    pub left_bound_inclusive: T,
405    pub right_bound_exclusive: T,
406}
407
408pub type CdcTableSnapshotSplit = CdcTableSnapshotSplitCommon<OwnedRow>;
409pub type CdcTableSnapshotSplitRaw = CdcTableSnapshotSplitCommon<Vec<u8>>;
410
411#[inline]
412pub fn build_pulsar_ack_channel_id(source_id: SourceId, split_id: &SplitId) -> String {
413    format!("{}-{}", source_id, split_id)
414}
risingwave_connector/source/mod.rs

risingwave_connector/source/
mod.rs
