risingwave_connector/source/
base.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
15use std::collections::{BTreeMap, HashMap};
16use std::sync::Arc;
17
18use anyhow::anyhow;
19use async_trait::async_trait;
20use bytes::Bytes;
21use enum_as_inner::EnumAsInner;
22use futures::future::try_join_all;
23use futures::stream::BoxStream;
24use futures::{Stream, StreamExt, TryStreamExt};
25use itertools::Itertools;
26use risingwave_common::array::StreamChunk;
27use risingwave_common::bail;
28use risingwave_common::id::{ActorId, FragmentId, SourceId};
29use risingwave_common::secret::LocalSecretManager;
30use risingwave_common::types::{JsonbVal, Scalar};
31use risingwave_pb::catalog::{PbSource, PbStreamSourceInfo};
32use risingwave_pb::plan_common::ExternalTableDesc;
33use risingwave_pb::source::ConnectorSplit;
34use rw_futures_util::select_all;
35use serde::de::DeserializeOwned;
36use serde_json::json;
37use tokio::sync::mpsc;
38
39use super::cdc::DebeziumCdcMeta;
40use super::datagen::DatagenMeta;
41use super::google_pubsub::GooglePubsubMeta;
42use super::kafka::KafkaMeta;
43use super::kinesis::KinesisMeta;
44use super::monitor::SourceMetrics;
45use super::nats::source::NatsMeta;
46use super::nexmark::source::message::NexmarkMeta;
47use super::pulsar::source::PulsarMeta;
48use crate::enforce_secret::EnforceSecret;
49use crate::error::ConnectorResult as Result;
50use crate::parser::ParserConfig;
51use crate::parser::schema_change::SchemaChangeEnvelope;
52use crate::source::SplitImpl::{
53    CitusCdc, MongodbCdc, MysqlCdc, OracleCdc, PostgresCdc, SqlServerCdc,
54};
55use crate::source::batch::BatchSourceSplitImpl;
56use crate::source::monitor::EnumeratorMetrics;
57use crate::with_options::WithOptions;
58use crate::{
59    WithOptionsSecResolved, WithPropertiesExt, dispatch_source_prop, dispatch_split_impl,
60    for_all_connections, for_all_sources, impl_connection, impl_connector_properties, impl_split,
61    match_source_name_str,
62};
63
64const SPLIT_TYPE_FIELD: &str = "split_type";
65const SPLIT_INFO_FIELD: &str = "split_info";
66pub const UPSTREAM_SOURCE_KEY: &str = "connector";
67
68pub const WEBHOOK_CONNECTOR: &str = "webhook";
69
70/// Callback wrapper for reporting CDC auto schema change fail events
71/// Parameters: (`table_id`, `table_name`, `cdc_table_id`, `upstream_ddl`, `fail_info`)
72#[derive(Clone)]
73pub struct CdcAutoSchemaChangeFailCallback(
74    Arc<dyn Fn(SourceId, String, String, String, String) + Send + Sync>,
75);
76
77impl CdcAutoSchemaChangeFailCallback {
78    pub fn new<F>(f: F) -> Self
79    where
80        F: Fn(SourceId, String, String, String, String) + Send + Sync + 'static,
81    {
82        Self(Arc::new(f))
83    }
84
85    pub fn call(
86        &self,
87        source_id: SourceId,
88        table_name: String,
89        cdc_table_id: String,
90        upstream_ddl: String,
91        fail_info: String,
92    ) {
93        self.0(source_id, table_name, cdc_table_id, upstream_ddl, fail_info);
94    }
95}
96
97impl std::fmt::Debug for CdcAutoSchemaChangeFailCallback {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        f.write_str("CdcAutoSchemaChangeFailCallback")
100    }
101}
102pub trait TryFromBTreeMap: Sized + UnknownFields {
103    /// Used to initialize the source properties from the raw untyped `WITH` options.
104    fn try_from_btreemap(
105        props: BTreeMap<String, String>,
106        deny_unknown_fields: bool,
107    ) -> Result<Self>;
108}
109
110/// Represents `WITH` options for sources.
111///
112/// Each instance should add a `#[derive(with_options::WithOptions)]` marker.
113pub trait SourceProperties:
114    TryFromBTreeMap + Clone + WithOptions + std::fmt::Debug + EnforceSecret
115{
116    const SOURCE_NAME: &'static str;
117    type Split: SplitMetaData
118        + TryFrom<SplitImpl, Error = crate::error::ConnectorError>
119        + Into<SplitImpl>;
120    type SplitEnumerator: SplitEnumerator<Properties = Self, Split = Self::Split>;
121    type SplitReader: SplitReader<Split = Self::Split, Properties = Self>;
122
123    /// Load additional info from `PbSource`. Currently only used by CDC.
124    fn init_from_pb_source(&mut self, _source: &PbSource) {}
125
126    /// Load additional info from `ExternalTableDesc`. Currently only used by CDC.
127    fn init_from_pb_cdc_table_desc(&mut self, _table_desc: &ExternalTableDesc) {}
128}
129
130pub trait UnknownFields {
131    /// Unrecognized fields in the `WITH` clause.
132    fn unknown_fields(&self) -> HashMap<String, String>;
133}
134
135impl<P: DeserializeOwned + UnknownFields> TryFromBTreeMap for P {
136    fn try_from_btreemap(
137        props: BTreeMap<String, String>,
138        deny_unknown_fields: bool,
139    ) -> Result<Self> {
140        let json_value = serde_json::to_value(props)?;
141        let res = serde_json::from_value::<P>(json_value)?;
142
143        if !deny_unknown_fields || res.unknown_fields().is_empty() {
144            Ok(res)
145        } else {
146            bail!(
147                "Unknown fields in the WITH clause: {:?}",
148                res.unknown_fields()
149            )
150        }
151    }
152}
153
154#[derive(Default)]
155pub struct CreateSplitReaderOpt {
156    pub support_multiple_splits: bool,
157    pub seek_to_latest: bool,
158}
159
160#[derive(Default)]
161pub struct CreateSplitReaderResult {
162    pub latest_splits: Option<Vec<SplitImpl>>,
163    pub backfill_info: HashMap<SplitId, BackfillInfo>,
164}
165
166pub async fn create_split_readers<P: SourceProperties>(
167    prop: P,
168    splits: Vec<SplitImpl>,
169    parser_config: ParserConfig,
170    source_ctx: SourceContextRef,
171    columns: Option<Vec<Column>>,
172    opt: CreateSplitReaderOpt,
173) -> Result<(BoxSourceReaderEventStream, CreateSplitReaderResult)> {
174    let splits = splits.into_iter().map(P::Split::try_from).try_collect()?;
175    let mut res = CreateSplitReaderResult {
176        backfill_info: HashMap::new(),
177        latest_splits: None,
178    };
179    if opt.support_multiple_splits {
180        let mut reader = P::SplitReader::new(
181            prop.clone(),
182            splits,
183            parser_config.clone(),
184            source_ctx.clone(),
185            columns.clone(),
186        )
187        .await?;
188        if opt.seek_to_latest {
189            res.latest_splits = Some(reader.seek_to_latest().await?);
190        }
191        res.backfill_info = reader.backfill_info();
192        Ok((reader.into_event_stream().boxed(), res))
193    } else {
194        let mut readers = try_join_all(splits.into_iter().map(|split| {
195            // TODO: is this reader split across multiple threads...? Realistically, we want
196            // source_ctx to live in a single actor.
197            P::SplitReader::new(
198                prop.clone(),
199                vec![split],
200                parser_config.clone(),
201                source_ctx.clone(),
202                columns.clone(),
203            )
204        }))
205        .await?;
206        if opt.seek_to_latest {
207            let mut latest_splits = vec![];
208            for reader in &mut readers {
209                latest_splits.extend(reader.seek_to_latest().await?);
210            }
211            res.latest_splits = Some(latest_splits);
212        }
213        res.backfill_info = readers.iter().flat_map(|r| r.backfill_info()).collect();
214        Ok((
215            select_all(readers.into_iter().map(|r| r.into_event_stream())).boxed(),
216            res,
217        ))
218    }
219}
220
221/// [`SplitEnumerator`] fetches the split metadata from the external source service.
222/// NOTE: It runs in the meta server, so probably it should be moved to the `meta` crate.
223#[async_trait]
224pub trait SplitEnumerator: Sized + Send {
225    type Split: SplitMetaData + Send;
226    type Properties;
227
228    async fn new(properties: Self::Properties, context: SourceEnumeratorContextRef)
229    -> Result<Self>;
230    async fn list_splits(&mut self) -> Result<Vec<Self::Split>>;
231    /// Do some cleanup work when a fragment is dropped, e.g., drop Kafka consumer group.
232    async fn on_drop_fragments(&mut self, _fragment_ids: Vec<FragmentId>) -> Result<()> {
233        Ok(())
234    }
235    /// Do some cleanup work when a backfill fragment is finished, e.g., drop Kafka consumer group.
236    async fn on_finish_backfill(&mut self, _fragment_ids: Vec<FragmentId>) -> Result<()> {
237        Ok(())
238    }
239    /// Called after `worker.tick()` execution to perform periodic operations,
240    /// such as monitoring upstream PostgreSQL `confirmed_flush_lsn`, etc.
241    /// This can be extended to support more periodic operations in the future.
242    async fn on_tick(&mut self) -> Result<()> {
243        Ok(())
244    }
245}
246
247pub type SourceContextRef = Arc<SourceContext>;
248pub type SourceEnumeratorContextRef = Arc<SourceEnumeratorContext>;
249
250/// Dyn-compatible [`SplitEnumerator`].
251#[async_trait]
252pub trait AnySplitEnumerator: Send {
253    async fn list_splits(&mut self) -> Result<Vec<SplitImpl>>;
254    async fn on_drop_fragments(&mut self, fragment_ids: Vec<FragmentId>) -> Result<()>;
255    async fn on_finish_backfill(&mut self, fragment_ids: Vec<FragmentId>) -> Result<()>;
256    async fn on_tick(&mut self) -> Result<()>;
257}
258
259#[async_trait]
260impl<T: SplitEnumerator<Split: Into<SplitImpl>> + 'static> AnySplitEnumerator for T {
261    async fn list_splits(&mut self) -> Result<Vec<SplitImpl>> {
262        SplitEnumerator::list_splits(self)
263            .await
264            .map(|s| s.into_iter().map(|s| s.into()).collect())
265    }
266
267    async fn on_drop_fragments(&mut self, fragment_ids: Vec<FragmentId>) -> Result<()> {
268        SplitEnumerator::on_drop_fragments(self, fragment_ids).await
269    }
270
271    async fn on_finish_backfill(&mut self, fragment_ids: Vec<FragmentId>) -> Result<()> {
272        SplitEnumerator::on_finish_backfill(self, fragment_ids).await
273    }
274
275    async fn on_tick(&mut self) -> Result<()> {
276        SplitEnumerator::on_tick(self).await
277    }
278}
279
280/// The max size of a chunk yielded by source stream.
281pub const MAX_CHUNK_SIZE: usize = 1024;
282
283#[derive(Debug, Clone, Copy)]
284pub struct SourceCtrlOpts {
285    /// The max size of a chunk yielded by source stream.
286    pub chunk_size: usize,
287    /// Whether to allow splitting a transaction into multiple chunks to meet the `max_chunk_size`.
288    pub split_txn: bool,
289}
290
291// The options in `SourceCtrlOpts` are so important that we don't want to impl `Default` for it,
292// so that we can prevent any unintentional use of the default value.
293impl !Default for SourceCtrlOpts {}
294
295impl SourceCtrlOpts {
296    #[cfg(test)]
297    pub fn for_test() -> Self {
298        SourceCtrlOpts {
299            chunk_size: 256,
300            split_txn: false,
301        }
302    }
303}
304
305#[derive(Debug)]
306pub struct SourceEnumeratorContext {
307    pub info: SourceEnumeratorInfo,
308    pub metrics: Arc<EnumeratorMetrics>,
309}
310
311impl SourceEnumeratorContext {
312    /// Create a dummy `SourceEnumeratorContext` for testing purpose, or for the situation
313    /// where the real context doesn't matter.
314    pub fn dummy() -> SourceEnumeratorContext {
315        SourceEnumeratorContext {
316            info: SourceEnumeratorInfo {
317                source_id: 0.into(),
318            },
319            metrics: Arc::new(EnumeratorMetrics::default()),
320        }
321    }
322}
323
324#[derive(Clone, Debug)]
325pub struct SourceEnumeratorInfo {
326    pub source_id: SourceId,
327}
328
329#[derive(Clone, Debug)]
330pub struct SourceContext {
331    pub actor_id: ActorId,
332    pub source_id: SourceId,
333    pub fragment_id: FragmentId,
334    pub source_name: String,
335    pub metrics: Arc<SourceMetrics>,
336    pub source_ctrl_opts: SourceCtrlOpts,
337    pub connector_props: ConnectorProperties,
338    // source parser put schema change event into this channel
339    pub schema_change_tx:
340        Option<mpsc::Sender<(SchemaChangeEnvelope, tokio::sync::oneshot::Sender<()>)>>,
341    // callback function to report CDC auto schema change fail events
342    pub on_cdc_auto_schema_change_failure: Option<CdcAutoSchemaChangeFailCallback>,
343}
344
345impl SourceContext {
346    pub fn new(
347        actor_id: ActorId,
348        source_id: SourceId,
349        fragment_id: FragmentId,
350        source_name: String,
351        metrics: Arc<SourceMetrics>,
352        source_ctrl_opts: SourceCtrlOpts,
353        connector_props: ConnectorProperties,
354        schema_change_channel: Option<
355            mpsc::Sender<(SchemaChangeEnvelope, tokio::sync::oneshot::Sender<()>)>,
356        >,
357    ) -> Self {
358        Self::new_with_auto_schema_change_callback(
359            actor_id,
360            source_id,
361            fragment_id,
362            source_name,
363            metrics,
364            source_ctrl_opts,
365            connector_props,
366            schema_change_channel,
367            None,
368        )
369    }
370
371    pub fn new_with_auto_schema_change_callback(
372        actor_id: ActorId,
373        source_id: SourceId,
374        fragment_id: FragmentId,
375        source_name: String,
376        metrics: Arc<SourceMetrics>,
377        source_ctrl_opts: SourceCtrlOpts,
378        connector_props: ConnectorProperties,
379        schema_change_channel: Option<
380            mpsc::Sender<(SchemaChangeEnvelope, tokio::sync::oneshot::Sender<()>)>,
381        >,
382        on_cdc_auto_schema_change_failure: Option<CdcAutoSchemaChangeFailCallback>,
383    ) -> Self {
384        Self {
385            actor_id,
386            source_id,
387            fragment_id,
388            source_name,
389            metrics,
390            source_ctrl_opts,
391            connector_props,
392            schema_change_tx: schema_change_channel,
393            on_cdc_auto_schema_change_failure,
394        }
395    }
396
397    /// Create a dummy `SourceContext` for testing purpose, or for the situation
398    /// where the real context doesn't matter.
399    pub fn dummy() -> Self {
400        Self::new(
401            0.into(),
402            SourceId::new(0),
403            0.into(),
404            "dummy".to_owned(),
405            Arc::new(SourceMetrics::default()),
406            SourceCtrlOpts {
407                chunk_size: MAX_CHUNK_SIZE,
408                split_txn: false,
409            },
410            ConnectorProperties::default(),
411            None,
412        )
413    }
414
415    /// Report CDC auto schema change fail event
416    /// Parameters: (`source_id`, `table_name`, `cdc_table_id`, `upstream_ddl`, `fail_info`)
417    pub fn on_cdc_auto_schema_change_failure(
418        &self,
419        source_id: SourceId,
420        table_name: String,
421        cdc_table_id: String,
422        upstream_ddl: String,
423        fail_info: String,
424    ) {
425        if let Some(ref cdc_auto_schema_change_fail_callback) =
426            self.on_cdc_auto_schema_change_failure
427        {
428            cdc_auto_schema_change_fail_callback.call(
429                source_id,
430                table_name,
431                cdc_table_id,
432                upstream_ddl,
433                fail_info,
434            );
435        }
436    }
437}
438
439#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
440pub enum SourceFormat {
441    #[default]
442    Invalid,
443    Native,
444    None,
445    Debezium,
446    DebeziumMongo,
447    Maxwell,
448    Canal,
449    Upsert,
450    Plain,
451}
452
453/// Refer to [`crate::parser::EncodingProperties`]
454#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
455pub enum SourceEncode {
456    #[default]
457    Invalid,
458    Native,
459    None,
460    Avro,
461    Csv,
462    Protobuf,
463    Json,
464    Bytes,
465    Parquet,
466}
467
468#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
469pub struct SourceStruct {
470    pub format: SourceFormat,
471    pub encode: SourceEncode,
472}
473
474impl SourceStruct {
475    pub fn new(format: SourceFormat, encode: SourceEncode) -> Self {
476        Self { format, encode }
477    }
478}
479
480// Only return valid (format, encode)
481pub fn extract_source_struct(info: &PbStreamSourceInfo) -> Result<SourceStruct> {
482    use risingwave_pb::plan_common::{PbEncodeType, PbFormatType, RowFormatType};
483
484    // old version meta.
485    if let Ok(format) = info.get_row_format() {
486        let (format, encode) = match format {
487            RowFormatType::Json => (SourceFormat::Plain, SourceEncode::Json),
488            RowFormatType::Protobuf => (SourceFormat::Plain, SourceEncode::Protobuf),
489            RowFormatType::DebeziumJson => (SourceFormat::Debezium, SourceEncode::Json),
490            RowFormatType::Avro => (SourceFormat::Plain, SourceEncode::Avro),
491            RowFormatType::Maxwell => (SourceFormat::Maxwell, SourceEncode::Json),
492            RowFormatType::CanalJson => (SourceFormat::Canal, SourceEncode::Json),
493            RowFormatType::Csv => (SourceFormat::Plain, SourceEncode::Csv),
494            RowFormatType::Native => (SourceFormat::Native, SourceEncode::Native),
495            RowFormatType::DebeziumAvro => (SourceFormat::Debezium, SourceEncode::Avro),
496            RowFormatType::UpsertJson => (SourceFormat::Upsert, SourceEncode::Json),
497            RowFormatType::UpsertAvro => (SourceFormat::Upsert, SourceEncode::Avro),
498            RowFormatType::DebeziumMongoJson => (SourceFormat::DebeziumMongo, SourceEncode::Json),
499            RowFormatType::Bytes => (SourceFormat::Plain, SourceEncode::Bytes),
500            RowFormatType::RowUnspecified => unreachable!(),
501        };
502        return Ok(SourceStruct::new(format, encode));
503    }
504    let source_format = info.get_format()?;
505    let source_encode = info.get_row_encode()?;
506    let (format, encode) = match (source_format, source_encode) {
507        (PbFormatType::Plain, PbEncodeType::Json) => (SourceFormat::Plain, SourceEncode::Json),
508        (PbFormatType::Plain, PbEncodeType::Protobuf) => {
509            (SourceFormat::Plain, SourceEncode::Protobuf)
510        }
511        (PbFormatType::Debezium, PbEncodeType::Json) => {
512            (SourceFormat::Debezium, SourceEncode::Json)
513        }
514        (PbFormatType::Plain, PbEncodeType::Avro) => (SourceFormat::Plain, SourceEncode::Avro),
515        (PbFormatType::Maxwell, PbEncodeType::Json) => (SourceFormat::Maxwell, SourceEncode::Json),
516        (PbFormatType::Canal, PbEncodeType::Json) => (SourceFormat::Canal, SourceEncode::Json),
517        (PbFormatType::Plain, PbEncodeType::Csv) => (SourceFormat::Plain, SourceEncode::Csv),
518        (PbFormatType::Plain, PbEncodeType::Parquet) => {
519            (SourceFormat::Plain, SourceEncode::Parquet)
520        }
521        (PbFormatType::Native, PbEncodeType::Native) => {
522            (SourceFormat::Native, SourceEncode::Native)
523        }
524        (PbFormatType::None, PbEncodeType::None) => (SourceFormat::None, SourceEncode::None),
525        (PbFormatType::Debezium, PbEncodeType::Avro) => {
526            (SourceFormat::Debezium, SourceEncode::Avro)
527        }
528        (PbFormatType::Upsert, PbEncodeType::Json) => (SourceFormat::Upsert, SourceEncode::Json),
529        (PbFormatType::Upsert, PbEncodeType::Avro) => (SourceFormat::Upsert, SourceEncode::Avro),
530        (PbFormatType::DebeziumMongo, PbEncodeType::Json) => {
531            (SourceFormat::DebeziumMongo, SourceEncode::Json)
532        }
533        (PbFormatType::Plain, PbEncodeType::Bytes) => (SourceFormat::Plain, SourceEncode::Bytes),
534        (PbFormatType::Upsert, PbEncodeType::Protobuf) => {
535            (SourceFormat::Upsert, SourceEncode::Protobuf)
536        }
537        (format, encode) => {
538            bail!(
539                "Unsupported combination of format {:?} and encode {:?}",
540                format,
541                encode
542            );
543        }
544    };
545    Ok(SourceStruct::new(format, encode))
546}
547
548/// Stream of [`SourceMessage`]. Messages flow through the stream in the unit of a batch.
549pub type BoxSourceMessageStream =
550    BoxStream<'static, crate::error::ConnectorResult<Vec<SourceMessage>>>;
551/// Stream of source message events.
552pub type BoxSourceMessageEventStream =
553    BoxStream<'static, crate::error::ConnectorResult<SourceMessageEvent>>;
554/// Stream of [`StreamChunk`]s parsed from the messages from the external source.
555pub type BoxSourceChunkStream = BoxStream<'static, crate::error::ConnectorResult<StreamChunk>>;
556/// Stream of source reader events.
557pub type BoxSourceReaderEventStream =
558    BoxStream<'static, crate::error::ConnectorResult<SourceReaderEvent>>;
559/// `StreamChunk` with the latest split state.
560/// The state is constructed in `StreamReaderBuilder::into_retry_stream`
561pub type StreamChunkWithState = (StreamChunk, HashMap<SplitId, SplitImpl>);
562/// See [`StreamChunkWithState`].
563pub type BoxSourceChunkWithStateStream =
564    BoxStream<'static, crate::error::ConnectorResult<StreamChunkWithState>>;
565
566#[derive(Debug)]
567pub enum SourceMessageEvent {
568    Data(Vec<SourceMessage>),
569    SplitProgress(HashMap<SplitId, String>),
570}
571
572#[derive(Debug)]
573pub enum SourceReaderEvent {
574    DataChunk(StreamChunk),
575    SplitProgress(HashMap<SplitId, String>),
576}
577
578/// Stream of [`Option<StreamChunk>`]s parsed from the messages from the external source.
579pub type BoxStreamingFileSourceChunkStream =
580    BoxStream<'static, crate::error::ConnectorResult<Option<StreamChunk>>>;
581
582// Manually expand the trait alias to improve IDE experience.
583pub trait SourceChunkStream:
584    Stream<Item = crate::error::ConnectorResult<StreamChunk>> + Send + 'static
585{
586}
587impl<T> SourceChunkStream for T where
588    T: Stream<Item = crate::error::ConnectorResult<StreamChunk>> + Send + 'static
589{
590}
591
592pub type BoxTryStream<M> = BoxStream<'static, crate::error::ConnectorResult<M>>;
593
594/// [`SplitReader`] is a new abstraction of the external connector read interface which is
595/// responsible for parsing, it is used to read messages from the outside and transform them into a
596/// stream of parsed [`StreamChunk`]
597#[async_trait]
598pub trait SplitReader: Sized + Send {
599    type Properties;
600    type Split: SplitMetaData;
601
602    async fn new(
603        properties: Self::Properties,
604        state: Vec<Self::Split>,
605        parser_config: ParserConfig,
606        source_ctx: SourceContextRef,
607        columns: Option<Vec<Column>>,
608    ) -> crate::error::ConnectorResult<Self>;
609
610    fn into_stream(self) -> BoxSourceChunkStream;
611
612    fn into_event_stream(self) -> BoxSourceReaderEventStream {
613        self.into_stream()
614            .map_ok(SourceReaderEvent::DataChunk)
615            .boxed()
616    }
617
618    fn backfill_info(&self) -> HashMap<SplitId, BackfillInfo> {
619        HashMap::new()
620    }
621
622    async fn seek_to_latest(&mut self) -> Result<Vec<SplitImpl>> {
623        Err(anyhow!("seek_to_latest is not supported for this connector").into())
624    }
625}
626
627/// Information used to determine whether we should start and finish source backfill.
628///
629/// XXX: if a connector cannot provide the latest offsets (but we want to make it shareable),
630/// perhaps we should ban blocking DDL for it.
631#[derive(Debug, Clone)]
632pub enum BackfillInfo {
633    HasDataToBackfill {
634        /// The last available offsets for each split (**inclusive**).
635        ///
636        /// This will be used to determine whether source backfill is finished when
637        /// there are no _new_ messages coming from upstream `SourceExecutor`. Otherwise,
638        /// blocking DDL cannot finish until new messages come.
639        ///
640        /// When there are upstream messages, we will use the latest offsets from the upstream.
641        latest_offset: String,
642    },
643    /// If there are no messages in the split at all, we don't need to start backfill.
644    /// In this case, there will be no message from the backfill stream too.
645    /// If we started backfill, we cannot finish it until new messages come.
646    /// So we mark this a special case for optimization.
647    NoDataToBackfill,
648}
649
650for_all_sources!(impl_connector_properties);
651
652impl Default for ConnectorProperties {
653    fn default() -> Self {
654        ConnectorProperties::Test(Box::default())
655    }
656}
657
658impl ConnectorProperties {
659    /// Creates typed source properties from the raw `WITH` properties.
660    ///
661    /// It checks the `connector` field, and them dispatches to the corresponding type's `try_from_btreemap` method.
662    ///
663    /// `deny_unknown_fields`: Since `WITH` options are persisted in meta, we do not deny unknown fields when restoring from
664    /// existing data to avoid breaking backwards compatibility. We only deny unknown fields when creating new sources.
665    pub fn extract(
666        with_properties: WithOptionsSecResolved,
667        deny_unknown_fields: bool,
668    ) -> Result<Self> {
669        let (options, secret_refs) = with_properties.into_parts();
670        let mut options_with_secret =
671            LocalSecretManager::global().fill_secrets(options, secret_refs)?;
672        let connector = options_with_secret
673            .remove(UPSTREAM_SOURCE_KEY)
674            .ok_or_else(|| anyhow!("Must specify 'connector' in WITH clause"))?
675            .to_lowercase();
676        match_source_name_str!(
677            connector.as_str(),
678            PropType,
679            PropType::try_from_btreemap(options_with_secret, deny_unknown_fields)
680                .map(ConnectorProperties::from),
681            |other| bail!("connector '{}' is not supported", other)
682        )
683    }
684
685    pub fn enforce_secret_source(
686        with_properties: &impl WithPropertiesExt,
687    ) -> crate::error::ConnectorResult<()> {
688        let connector = with_properties
689            .get_connector()
690            .ok_or_else(|| anyhow!("Must specify 'connector' in WITH clause"))?
691            .to_lowercase();
692        let key_iter = with_properties.key_iter();
693        match_source_name_str!(
694            connector.as_str(),
695            PropType,
696            PropType::enforce_secret(key_iter),
697            |other| bail!("connector '{}' is not supported", other)
698        )
699    }
700
701    pub fn enable_drop_split(&self) -> bool {
702        // enable split scale in just for Kinesis
703        matches!(
704            self,
705            ConnectorProperties::Kinesis(_)
706                | ConnectorProperties::Nats(_)
707                | ConnectorProperties::GooglePubsub(_)
708        )
709    }
710
711    /// For most connectors, this should be false. When enabled, RisingWave should not track any progress.
712    pub fn enable_adaptive_splits(&self) -> bool {
713        matches!(
714            self,
715            ConnectorProperties::Nats(_) | ConnectorProperties::GooglePubsub(_)
716        )
717    }
718
719    /// Load additional info from `PbSource`. Currently only used by CDC.
720    pub fn init_from_pb_source(&mut self, source: &PbSource) {
721        dispatch_source_prop!(self, |prop| prop.init_from_pb_source(source))
722    }
723
724    /// Load additional info from `ExternalTableDesc`. Currently only used by CDC.
725    pub fn init_from_pb_cdc_table_desc(&mut self, cdc_table_desc: &ExternalTableDesc) {
726        dispatch_source_prop!(self, |prop| prop
727            .init_from_pb_cdc_table_desc(cdc_table_desc))
728    }
729
730    pub fn support_multiple_splits(&self) -> bool {
731        matches!(self, ConnectorProperties::Kafka(_))
732            || matches!(self, ConnectorProperties::OpendalS3(_))
733            || matches!(self, ConnectorProperties::Gcs(_))
734            || matches!(self, ConnectorProperties::Azblob(_))
735    }
736
737    pub async fn create_split_enumerator(
738        self,
739        context: crate::source::base::SourceEnumeratorContextRef,
740    ) -> crate::error::ConnectorResult<Box<dyn AnySplitEnumerator>> {
741        let enumerator: Box<dyn AnySplitEnumerator> = dispatch_source_prop!(self, |prop| Box::new(
742            <PropType as SourceProperties>::SplitEnumerator::new(*prop, context).await?
743        ));
744        Ok(enumerator)
745    }
746
747    pub async fn create_split_reader(
748        self,
749        splits: Vec<SplitImpl>,
750        parser_config: ParserConfig,
751        source_ctx: SourceContextRef,
752        columns: Option<Vec<Column>>,
753        mut opt: crate::source::CreateSplitReaderOpt,
754    ) -> Result<(
755        BoxSourceReaderEventStream,
756        crate::source::CreateSplitReaderResult,
757    )> {
758        opt.support_multiple_splits = self.support_multiple_splits();
759        tracing::debug!(
760            ?splits,
761            support_multiple_splits = opt.support_multiple_splits,
762            "spawning connector split reader",
763        );
764
765        dispatch_source_prop!(self, |prop| create_split_readers(
766            *prop,
767            splits,
768            parser_config,
769            source_ctx,
770            columns,
771            opt
772        )
773        .await)
774    }
775}
776
777for_all_sources!(impl_split);
778for_all_connections!(impl_connection);
779
780impl From<&SplitImpl> for ConnectorSplit {
781    fn from(split: &SplitImpl) -> Self {
782        dispatch_split_impl!(split, |inner| {
783            ConnectorSplit {
784                split_type: String::from(PropType::SOURCE_NAME),
785                encoded_split: inner.encode_to_bytes().to_vec(),
786            }
787        })
788    }
789}
790
791impl TryFrom<&ConnectorSplit> for SplitImpl {
792    type Error = crate::error::ConnectorError;
793
794    fn try_from(split: &ConnectorSplit) -> std::result::Result<Self, Self::Error> {
795        let split_type = split.split_type.to_lowercase();
796        match_source_name_str!(
797            split_type.as_str(),
798            PropType,
799            {
800                <PropType as SourceProperties>::Split::restore_from_bytes(
801                    split.encoded_split.as_ref(),
802                )
803                .map(Into::into)
804            },
805            |other| bail!("connector '{}' is not supported", other)
806        )
807    }
808}
809
810impl SplitImpl {
811    fn restore_from_json_inner(split_type: &str, value: JsonbVal) -> Result<Self> {
812        let split_type = split_type.to_lowercase();
813        match_source_name_str!(
814            split_type.as_str(),
815            PropType,
816            <PropType as SourceProperties>::Split::restore_from_json(value).map(Into::into),
817            |other| bail!("connector '{}' is not supported", other)
818        )
819    }
820
821    pub fn is_cdc_split(&self) -> bool {
822        matches!(
823            self,
824            MysqlCdc(_)
825                | PostgresCdc(_)
826                | MongodbCdc(_)
827                | CitusCdc(_)
828                | SqlServerCdc(_)
829                | OracleCdc(_)
830        )
831    }
832
833    /// Get the current split offset.
834    pub fn get_cdc_split_offset(&self) -> String {
835        match self {
836            MysqlCdc(split) => split.start_offset().clone().unwrap_or_default(),
837            PostgresCdc(split) => split.start_offset().clone().unwrap_or_default(),
838            MongodbCdc(split) => split.start_offset().clone().unwrap_or_default(),
839            CitusCdc(split) => split.start_offset().clone().unwrap_or_default(),
840            SqlServerCdc(split) => split.start_offset().clone().unwrap_or_default(),
841            OracleCdc(split) => split.start_offset().clone().unwrap_or_default(),
842            _ => unreachable!("get_cdc_split_offset() is only for cdc split"),
843        }
844    }
845
846    pub fn into_batch_split(self) -> Option<BatchSourceSplitImpl> {
847        match self {
848            SplitImpl::BatchPosixFs(batch_posix_fs_split) => {
849                Some(BatchSourceSplitImpl::BatchPosixFs(batch_posix_fs_split))
850            }
851            _ => None,
852        }
853    }
854}
855
856impl SplitMetaData for SplitImpl {
857    fn id(&self) -> SplitId {
858        dispatch_split_impl!(self, |inner| inner.id())
859    }
860
861    fn encode_to_json(&self) -> JsonbVal {
862        use serde_json::json;
863        let inner = self.encode_to_json_inner().take();
864        json!({ SPLIT_TYPE_FIELD: self.get_type(), SPLIT_INFO_FIELD: inner}).into()
865    }
866
867    fn restore_from_json(value: JsonbVal) -> Result<Self> {
868        let mut value = value.take();
869        let json_obj = value.as_object_mut().unwrap();
870        let split_type = json_obj
871            .remove(SPLIT_TYPE_FIELD)
872            .unwrap()
873            .as_str()
874            .unwrap()
875            .to_owned();
876        let inner_value = json_obj.remove(SPLIT_INFO_FIELD).unwrap();
877        Self::restore_from_json_inner(&split_type, inner_value.into())
878    }
879
880    fn update_offset(&mut self, last_seen_offset: String) -> Result<()> {
881        dispatch_split_impl!(self, |inner| inner.update_offset(last_seen_offset))
882    }
883}
884
885impl SplitImpl {
886    pub fn get_type(&self) -> String {
887        dispatch_split_impl!(self, |_inner| PropType::SOURCE_NAME.to_owned())
888    }
889
890    pub fn update_in_place(&mut self, last_seen_offset: String) -> Result<()> {
891        dispatch_split_impl!(self, |inner| inner.update_offset(last_seen_offset)?);
892        Ok(())
893    }
894
895    pub fn encode_to_json_inner(&self) -> JsonbVal {
896        dispatch_split_impl!(self, |inner| inner.encode_to_json())
897    }
898}
899
900use risingwave_common::types::DataType;
901
902#[derive(Clone, Debug)]
903pub struct Column {
904    pub name: String,
905    pub data_type: DataType,
906    /// This field is only used by datagen.
907    pub is_visible: bool,
908}
909
910/// Split id resides in every source message, use `Arc` to avoid copying.
911pub type SplitId = Arc<str>;
912
913/// The message pumped from the external source service.
914/// The third-party message structs will eventually be transformed into this struct.
915#[derive(Debug, Clone)]
916pub struct SourceMessage {
917    pub key: Option<Vec<u8>>,
918    pub payload: Option<Vec<u8>>,
919    pub offset: String, // TODO: use `Arc<str>`
920    pub split_id: SplitId,
921    pub meta: SourceMeta,
922}
923
924impl SourceMessage {
925    /// Create a dummy `SourceMessage` with all fields unset for testing purposes.
926    pub fn dummy() -> Self {
927        Self {
928            key: None,
929            payload: None,
930            offset: "".to_owned(),
931            split_id: "".into(),
932            meta: SourceMeta::Empty,
933        }
934    }
935
936    /// Check whether the source message is a CDC heartbeat message.
937    pub fn is_cdc_heartbeat(&self) -> bool {
938        self.key.is_none() && self.payload.is_none()
939    }
940}
941
942#[derive(Debug, Clone)]
943pub enum SourceMeta {
944    Kafka(KafkaMeta),
945    Kinesis(KinesisMeta),
946    Pulsar(PulsarMeta),
947    Nexmark(NexmarkMeta),
948    GooglePubsub(GooglePubsubMeta),
949    Datagen(DatagenMeta),
950    DebeziumCdc(DebeziumCdcMeta),
951    Nats(NatsMeta),
952    // For the source that doesn't have meta data.
953    Empty,
954}
955
956/// Implement Eq manually to ignore the `meta` field.
957impl PartialEq for SourceMessage {
958    fn eq(&self, other: &Self) -> bool {
959        self.offset == other.offset
960            && self.split_id == other.split_id
961            && self.payload == other.payload
962    }
963}
964impl Eq for SourceMessage {}
965
966/// The metadata of a split.
967pub trait SplitMetaData: Sized {
968    fn id(&self) -> SplitId;
969    fn encode_to_bytes(&self) -> Bytes {
970        self.encode_to_json()
971            .as_scalar_ref()
972            .value_serialize()
973            .into()
974    }
975    fn restore_from_bytes(bytes: &[u8]) -> Result<Self> {
976        Self::restore_from_json(JsonbVal::value_deserialize(bytes).unwrap())
977    }
978
979    /// Encode the whole split metadata to a JSON object
980    fn encode_to_json(&self) -> JsonbVal;
981    fn restore_from_json(value: JsonbVal) -> Result<Self>;
982    fn update_offset(&mut self, last_seen_offset: String) -> crate::error::ConnectorResult<()>;
983}
984
985/// [`ConnectorState`] maintains the consuming splits' info. In specific split readers,
986/// `ConnectorState` cannot be [`None`] and contains one(for mq split readers) or many(for fs
987/// split readers) [`SplitImpl`]. If no split is assigned to source executor, `ConnectorState` is
988/// [`None`] and the created source stream will be a pending stream.
989pub type ConnectorState = Option<Vec<SplitImpl>>;
990
991#[cfg(test)]
992mod tests {
993    use maplit::*;
994    use nexmark::event::EventType;
995
996    use super::*;
997    use crate::source::cdc::{DebeziumCdcSplit, Mysql};
998    use crate::source::kafka::KafkaSplit;
999
1000    #[test]
1001    fn test_split_impl_get_fn() -> Result<()> {
1002        let split = KafkaSplit::new(0, Some(0), Some(0), "demo".to_owned());
1003        let split_impl = SplitImpl::Kafka(split.clone());
1004        let get_value = split_impl.into_kafka().unwrap();
1005        println!("{:?}", get_value);
1006        assert_eq!(split.encode_to_bytes(), get_value.encode_to_bytes());
1007        assert_eq!(split.encode_to_json(), get_value.encode_to_json());
1008
1009        Ok(())
1010    }
1011
1012    #[test]
1013    fn test_cdc_split_state() -> Result<()> {
1014        let offset_str = "{\"sourcePartition\":{\"server\":\"RW_CDC_mydb.products\"},\"sourceOffset\":{\"transaction_id\":null,\"ts_sec\":1670407377,\"file\":\"binlog.000001\",\"pos\":98587,\"row\":2,\"server_id\":1,\"event\":2}}";
1015        let split = DebeziumCdcSplit::<Mysql>::new(1001, Some(offset_str.to_owned()), None);
1016        let split_impl = SplitImpl::MysqlCdc(split);
1017        let encoded_split = split_impl.encode_to_bytes();
1018        let restored_split_impl = SplitImpl::restore_from_bytes(encoded_split.as_ref())?;
1019        assert_eq!(
1020            split_impl.encode_to_bytes(),
1021            restored_split_impl.encode_to_bytes()
1022        );
1023        assert_eq!(
1024            split_impl.encode_to_json(),
1025            restored_split_impl.encode_to_json()
1026        );
1027
1028        let encoded_split = split_impl.encode_to_json();
1029        let restored_split_impl = SplitImpl::restore_from_json(encoded_split)?;
1030        assert_eq!(
1031            split_impl.encode_to_bytes(),
1032            restored_split_impl.encode_to_bytes()
1033        );
1034        assert_eq!(
1035            split_impl.encode_to_json(),
1036            restored_split_impl.encode_to_json()
1037        );
1038        Ok(())
1039    }
1040
1041    #[test]
1042    fn test_extract_nexmark_config() {
1043        let props = convert_args!(btreemap!(
1044            "connector" => "nexmark",
1045            "nexmark.table.type" => "Person",
1046            "nexmark.split.num" => "1",
1047        ));
1048
1049        let props =
1050            ConnectorProperties::extract(WithOptionsSecResolved::without_secrets(props), true)
1051                .unwrap();
1052
1053        if let ConnectorProperties::Nexmark(props) = props {
1054            assert_eq!(props.table_type, Some(EventType::Person));
1055            assert_eq!(props.split_num, 1);
1056        } else {
1057            panic!("extract nexmark config failed");
1058        }
1059    }
1060
1061    #[test]
1062    fn test_extract_kafka_config() {
1063        let props = convert_args!(btreemap!(
1064            "connector" => "kafka",
1065            "properties.bootstrap.server" => "b1,b2",
1066            "topic" => "test",
1067            "scan.startup.mode" => "earliest",
1068            "broker.rewrite.endpoints" => r#"{"b-1:9092":"dns-1", "b-2:9092":"dns-2"}"#,
1069        ));
1070
1071        let props =
1072            ConnectorProperties::extract(WithOptionsSecResolved::without_secrets(props), true)
1073                .unwrap();
1074        if let ConnectorProperties::Kafka(k) = props {
1075            let btreemap = btreemap! {
1076                "b-1:9092".to_owned() => "dns-1".to_owned(),
1077                "b-2:9092".to_owned() => "dns-2".to_owned(),
1078            };
1079            assert_eq!(k.privatelink_common.broker_rewrite_map, Some(btreemap));
1080        } else {
1081            panic!("extract kafka config failed");
1082        }
1083    }
1084
1085    #[test]
1086    fn test_extract_cdc_properties() {
1087        let user_props_mysql = convert_args!(btreemap!(
1088            "connector" => "mysql-cdc",
1089            "database.hostname" => "127.0.0.1",
1090            "database.port" => "3306",
1091            "database.user" => "root",
1092            "database.password" => "123456",
1093            "database.name" => "mydb",
1094            "table.name" => "products",
1095        ));
1096
1097        let user_props_postgres = convert_args!(btreemap!(
1098            "connector" => "postgres-cdc",
1099            "database.hostname" => "127.0.0.1",
1100            "database.port" => "5432",
1101            "database.user" => "root",
1102            "database.password" => "654321",
1103            "schema.name" => "public",
1104            "database.name" => "mypgdb",
1105            "table.name" => "orders",
1106        ));
1107
1108        let conn_props = ConnectorProperties::extract(
1109            WithOptionsSecResolved::without_secrets(user_props_mysql),
1110            true,
1111        )
1112        .unwrap();
1113        if let ConnectorProperties::MysqlCdc(c) = conn_props {
1114            assert_eq!(c.properties.get("database.hostname").unwrap(), "127.0.0.1");
1115            assert_eq!(c.properties.get("database.port").unwrap(), "3306");
1116            assert_eq!(c.properties.get("database.user").unwrap(), "root");
1117            assert_eq!(c.properties.get("database.password").unwrap(), "123456");
1118            assert_eq!(c.properties.get("database.name").unwrap(), "mydb");
1119            assert_eq!(c.properties.get("table.name").unwrap(), "products");
1120        } else {
1121            panic!("extract cdc config failed");
1122        }
1123
1124        let conn_props = ConnectorProperties::extract(
1125            WithOptionsSecResolved::without_secrets(user_props_postgres),
1126            true,
1127        )
1128        .unwrap();
1129        if let ConnectorProperties::PostgresCdc(c) = conn_props {
1130            assert_eq!(c.properties.get("database.hostname").unwrap(), "127.0.0.1");
1131            assert_eq!(c.properties.get("database.port").unwrap(), "5432");
1132            assert_eq!(c.properties.get("database.user").unwrap(), "root");
1133            assert_eq!(c.properties.get("database.password").unwrap(), "654321");
1134            assert_eq!(c.properties.get("schema.name").unwrap(), "public");
1135            assert_eq!(c.properties.get("database.name").unwrap(), "mypgdb");
1136            assert_eq!(c.properties.get("table.name").unwrap(), "orders");
1137        } else {
1138            panic!("extract cdc config failed");
1139        }
1140    }
1141}
risingwave_connector/source/base.rs

risingwave_connector/source/
base.rs
