unify entry point for reader

Signed-off-by: xxchan <[email protected]>
risingwavelabs · Jan 10, 2025 · 11fb96f · 11fb96f
1 parent 5bde2a9
commit 11fb96f
Show file tree

Hide file tree

Showing 27 changed files with 153 additions and 175 deletions.
diff --git a/src/bench/sink_bench/main.rs b/src/bench/sink_bench/main.rs
@@ -59,7 +59,7 @@ use risingwave_connector::source::datagen::{
     DatagenProperties, DatagenSplitEnumerator, DatagenSplitReader,
 };
 use risingwave_connector::source::{
-    Column, DataType, SourceContext, SourceEnumeratorContext, SplitEnumerator, SplitReader,
+    Column, DataType, SourceContext, SourceEnumeratorContext, SplitReader,
 };
 use risingwave_stream::executor::test_utils::prelude::ColumnDesc;
 use risingwave_stream::executor::{Barrier, Message, MessageStreamItem, StreamExecutorError};
@@ -289,7 +289,6 @@ impl MockDatagenSource {
             properties.clone(),
             SourceEnumeratorContext::dummy().into(),
         )
-        .await
         .unwrap();
         let parser_config = ParserConfig {
             specific: SpecificParserConfig {

diff --git a/src/connector/src/macros.rs b/src/connector/src/macros.rs
@@ -328,6 +328,31 @@ macro_rules! impl_connector_properties {
                     };
                     Ok(enumerator)
                 }
+
+
+                pub async fn create_split_reader(
+                    self,
+                    splits: Vec<SplitImpl>,
+                    parser_config: ParserConfig,
+                    source_ctx: SourceContextRef,
+                    columns: Option<Vec<Column>>,
+                    mut opt: $crate::source::CreateSplitReaderOpt,
+                ) -> Result<(BoxSourceChunkStream, $crate::source::CreateSplitReaderResult)> {
+                    opt.support_multiple_splits = self.support_multiple_splits();
+                    tracing::debug!(
+                        ?splits,
+                        support_multiple_splits = opt.support_multiple_splits,
+                        "spawning connector split reader",
+                    );
+
+                    match self {
+                        $(
+                            ConnectorProperties::$variant_name(prop) => {
+                                $crate::source::create_split_readers(*prop, splits, parser_config, source_ctx, columns, opt).await
+                            }
+                        )*
+                    }
+                }
             }
         }
 

diff --git a/src/connector/src/source/base.rs b/src/connector/src/source/base.rs
@@ -20,8 +20,9 @@ use async_trait::async_trait;
 use aws_sdk_s3::types::Object;
 use bytes::Bytes;
 use enum_as_inner::EnumAsInner;
+use futures::future::try_join_all;
 use futures::stream::BoxStream;
-use futures::Stream;
+use futures::{Stream, StreamExt};
 use itertools::Itertools;
 use risingwave_common::array::StreamChunk;
 use risingwave_common::bail;
@@ -31,6 +32,7 @@ use risingwave_common::types::{JsonbVal, Scalar};
 use risingwave_pb::catalog::{PbSource, PbStreamSourceInfo};
 use risingwave_pb::plan_common::ExternalTableDesc;
 use risingwave_pb::source::ConnectorSplit;
+use rw_futures_util::select_all;
 use serde::de::DeserializeOwned;
 use serde_json::json;
 use tokio::sync::mpsc;
@@ -118,15 +120,71 @@ impl<P: DeserializeOwned + UnknownFields> TryFromBTreeMap for P {
     }
 }
 
-pub async fn create_split_reader<P: SourceProperties>(
+#[derive(Default)]
+pub struct CreateSplitReaderOpt {
+    pub support_multiple_splits: bool,
+    pub seek_to_latest: bool,
+}
+
+#[derive(Default)]
+pub struct CreateSplitReaderResult {
+    pub latest_splits: Option<Vec<SplitImpl>>,
+    pub backfill_info: HashMap<SplitId, BackfillInfo>,
+}
+
+pub async fn create_split_readers<P: SourceProperties>(
     prop: P,
     splits: Vec<SplitImpl>,
     parser_config: ParserConfig,
     source_ctx: SourceContextRef,
     columns: Option<Vec<Column>>,
-) -> Result<P::SplitReader> {
+    opt: CreateSplitReaderOpt,
+) -> Result<(BoxSourceChunkStream, CreateSplitReaderResult)> {
     let splits = splits.into_iter().map(P::Split::try_from).try_collect()?;
-    P::SplitReader::new(prop, splits, parser_config, source_ctx, columns).await
+    let mut res = CreateSplitReaderResult {
+        backfill_info: HashMap::new(),
+        latest_splits: None,
+    };
+    if opt.support_multiple_splits {
+        let mut reader = P::SplitReader::new(
+            prop.clone(),
+            splits,
+            parser_config.clone(),
+            source_ctx.clone(),
+            columns.clone(),
+        )
+        .await?;
+        if opt.seek_to_latest {
+            res.latest_splits = Some(reader.seek_to_latest().await?);
+        }
+        res.backfill_info = reader.backfill_info();
+        Ok((reader.into_stream().boxed(), res))
+    } else {
+        let mut readers = try_join_all(splits.into_iter().map(|split| {
+            // TODO: is this reader split across multiple threads...? Realistically, we want
+            // source_ctx to live in a single actor.
+            P::SplitReader::new(
+                prop.clone(),
+                vec![split],
+                parser_config.clone(),
+                source_ctx.clone(),
+                columns.clone(),
+            )
+        }))
+        .await?;
+        if opt.seek_to_latest {
+            let mut latest_splits = vec![];
+            for reader in &mut readers {
+                latest_splits.extend(reader.seek_to_latest().await?);
+            }
+            res.latest_splits = Some(latest_splits);
+        }
+        res.backfill_info = readers.iter().flat_map(|r| r.backfill_info()).collect();
+        Ok((
+            select_all(readers.into_iter().map(|r| r.into_stream())).boxed(),
+            res,
+        ))
+    }
 }
 
 /// [`SplitEnumerator`] fetches the split metadata from the external source service.

diff --git a/src/connector/src/source/cdc/enumerator/mod.rs b/src/connector/src/source/cdc/enumerator/mod.rs
@@ -30,7 +30,7 @@ use crate::source::cdc::{
     table_schema_exclude_additional_columns, CdcProperties, CdcSourceTypeTrait, Citus,
     DebeziumCdcSplit, Mongodb, Mysql, Postgres, SqlServer,
 };
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 pub const DATABASE_SERVERS_KEY: &str = "database.servers";
 

diff --git a/src/connector/src/source/cdc/mod.rs b/src/connector/src/source/cdc/mod.rs
@@ -213,7 +213,7 @@ where
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        Ok(DebeziumSplitEnumerator::new(self, context).await?)
+        DebeziumSplitEnumerator::new(self, context).await
     }
 }
 

diff --git a/src/connector/src/source/datagen/enumerator/mod.rs b/src/connector/src/source/datagen/enumerator/mod.rs
@@ -16,15 +16,15 @@ use anyhow::Context;
 use async_trait::async_trait;
 
 use crate::source::datagen::{DatagenProperties, DatagenSplit};
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct DatagenSplitEnumerator {
     split_num: i32,
 }
 
 impl DatagenSplitEnumerator {
-    pub async fn new(
+    pub fn new(
         properties: DatagenProperties,
         _context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<DatagenSplitEnumerator> {

diff --git a/src/connector/src/source/datagen/mod.rs b/src/connector/src/source/datagen/mod.rs
@@ -68,7 +68,7 @@ impl SourceProperties for DatagenProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        DatagenSplitEnumerator::new(self, context).await
+        DatagenSplitEnumerator::new(self, context)
     }
 }
 

diff --git a/src/connector/src/source/filesystem/opendal_source/mod.rs b/src/connector/src/source/filesystem/opendal_source/mod.rs
@@ -91,7 +91,7 @@ impl SourceProperties for GcsProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        OpendalEnumerator::new(self, context).await
+        OpendalEnumerator::new(self, context)
     }
 }
 
@@ -167,7 +167,7 @@ impl SourceProperties for OpendalS3Properties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        OpendalEnumerator::new(self, context).await
+        OpendalEnumerator::new(self, context)
     }
 }
 
@@ -207,7 +207,7 @@ impl SourceProperties for PosixFsProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        OpendalEnumerator::new(self, context).await
+        OpendalEnumerator::new(self, context)
     }
 }
 
@@ -253,7 +253,7 @@ impl SourceProperties for AzblobProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        OpendalEnumerator::new(self, context).await
+        OpendalEnumerator::new(self, context)
     }
 }
 

diff --git a/src/connector/src/source/filesystem/opendal_source/opendal_enumerator.rs b/src/connector/src/source/filesystem/opendal_source/opendal_enumerator.rs
@@ -26,7 +26,7 @@ use super::OpendalSource;
 use crate::error::ConnectorResult;
 use crate::source::filesystem::file_common::CompressionFormat;
 use crate::source::filesystem::{FsPageItem, OpendalFsSplit};
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 #[derive(Debug, Clone)]
 pub struct OpendalEnumerator<Src: OpendalSource> {
@@ -39,7 +39,7 @@ pub struct OpendalEnumerator<Src: OpendalSource> {
 }
 
 impl<Src: OpendalSource> OpendalEnumerator<Src> {
-    pub async fn new(
+    pub fn new(
         properties: Src::Properties,
         _context: crate::source::SourceEnumeratorContextRef,
     ) -> ConnectorResult<Self> {

diff --git a/src/connector/src/source/filesystem/s3/enumerator.rs b/src/connector/src/source/filesystem/s3/enumerator.rs
@@ -20,7 +20,7 @@ use crate::aws_utils::{default_conn_config, s3_client};
 use crate::connector_common::AwsAuthProps;
 use crate::source::filesystem::file_common::FsSplit;
 use crate::source::filesystem::s3::S3Properties;
-use crate::source::{FsListInner, SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::{FsListInner, SplitEnumerator};
 
 /// Get the prefix from a glob
 pub fn get_prefix(glob: &str) -> String {

diff --git a/src/connector/src/source/filesystem/s3/mod.rs b/src/connector/src/source/filesystem/s3/mod.rs
@@ -75,7 +75,7 @@ impl SourceProperties for S3Properties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        Ok(S3SplitEnumerator::new(self, context).await?)
+        S3SplitEnumerator::new(self, context).await
     }
 }
 

diff --git a/src/connector/src/source/google_pubsub/enumerator/client.rs b/src/connector/src/source/google_pubsub/enumerator/client.rs
@@ -22,7 +22,6 @@ use crate::error::ConnectorResult;
 use crate::source::base::SplitEnumerator;
 use crate::source::google_pubsub::split::PubsubSplit;
 use crate::source::google_pubsub::PubsubProperties;
-use crate::source::SourceEnumeratorContextRef;
 
 pub struct PubsubSplitEnumerator {
     subscription: String,

diff --git a/src/connector/src/source/iceberg/mod.rs b/src/connector/src/source/iceberg/mod.rs
@@ -41,8 +41,8 @@ use crate::connector_common::IcebergCommon;
 use crate::error::{ConnectorError, ConnectorResult};
 use crate::parser::ParserConfig;
 use crate::source::{
-    BoxSourceChunkStream, Column, SourceContextRef, SourceEnumeratorContextRef, SourceProperties,
-    SplitEnumerator, SplitId, SplitMetaData, SplitReader, UnknownFields,
+    BoxSourceChunkStream, Column, SourceContextRef, SourceProperties, SplitEnumerator, SplitId,
+    SplitMetaData, SplitReader, UnknownFields,
 };
 pub const ICEBERG_CONNECTOR: &str = "iceberg";
 
@@ -115,7 +115,7 @@ impl SourceProperties for IcebergProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        IcebergSplitEnumerator::new(self, context).await
+        IcebergSplitEnumerator::new(self, context)
     }
 }
 
@@ -256,7 +256,7 @@ pub struct IcebergSplitEnumerator {
 }
 
 impl IcebergSplitEnumerator {
-    pub async fn new(
+    pub fn new(
         properties: IcebergProperties,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> ConnectorResult<Self> {

diff --git a/src/connector/src/source/kafka/enumerator/client.rs b/src/connector/src/source/kafka/enumerator/client.rs
@@ -33,7 +33,6 @@ use crate::source::kafka::{
     KafkaConnectionProps, KafkaContextCommon, KafkaProperties, RwConsumerContext,
     KAFKA_ISOLATION_LEVEL,
 };
-use crate::source::SourceEnumeratorContextRef;
 
 type KafkaClientType = BaseConsumer<RwConsumerContext>;
 

diff --git a/src/connector/src/source/kinesis/enumerator/client.rs b/src/connector/src/source/kinesis/enumerator/client.rs
@@ -21,7 +21,7 @@ use risingwave_common::bail;
 use crate::error::ConnectorResult as Result;
 use crate::source::kinesis::split::KinesisOffset;
 use crate::source::kinesis::*;
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 pub struct KinesisSplitEnumerator {
     stream_name: String,

diff --git a/src/connector/src/source/mqtt/enumerator/mod.rs b/src/connector/src/source/mqtt/enumerator/mod.rs
@@ -26,7 +26,7 @@ use tokio::sync::RwLock;
 use super::source::MqttSplit;
 use super::MqttProperties;
 use crate::error::ConnectorResult;
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 pub struct MqttSplitEnumerator {
     #[expect(dead_code)]

diff --git a/src/connector/src/source/nats/enumerator/mod.rs b/src/connector/src/source/nats/enumerator/mod.rs
@@ -20,7 +20,7 @@ use risingwave_common::bail;
 use super::source::{NatsOffset, NatsSplit};
 use super::NatsProperties;
 use crate::error::ConnectorResult;
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator, SplitId};
+use crate::source::{SplitEnumerator, SplitId};
 
 #[derive(Debug, Clone)]
 pub struct NatsSplitEnumerator {

diff --git a/src/connector/src/source/nexmark/enumerator/mod.rs b/src/connector/src/source/nexmark/enumerator/mod.rs
@@ -17,16 +17,14 @@ use async_trait::async_trait;
 use crate::error::ConnectorResult;
 use crate::source::nexmark::split::NexmarkSplit;
 use crate::source::nexmark::NexmarkProperties;
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 pub struct NexmarkSplitEnumerator {
     split_num: i32,
 }
 
-impl NexmarkSplitEnumerator {}
-
 impl NexmarkSplitEnumerator {
-    pub async fn new(
+    pub fn new(
         properties: NexmarkProperties,
         _context: crate::source::SourceEnumeratorContextRef,
     ) -> ConnectorResult<NexmarkSplitEnumerator> {

diff --git a/src/connector/src/source/nexmark/mod.rs b/src/connector/src/source/nexmark/mod.rs
@@ -233,7 +233,7 @@ impl SourceProperties for NexmarkProperties {
         self,
         context: crate::source::SourceEnumeratorContextRef,
     ) -> crate::error::ConnectorResult<Self::SplitEnumerator> {
-        NexmarkSplitEnumerator::new(self, context).await
+        NexmarkSplitEnumerator::new(self, context)
     }
 }
 

diff --git a/src/connector/src/source/pulsar/enumerator/client.rs b/src/connector/src/source/pulsar/enumerator/client.rs
@@ -23,7 +23,7 @@ use crate::error::ConnectorResult;
 use crate::source::pulsar::split::PulsarSplit;
 use crate::source::pulsar::topic::{parse_topic, Topic};
 use crate::source::pulsar::PulsarProperties;
-use crate::source::{SourceEnumeratorContextRef, SplitEnumerator};
+use crate::source::SplitEnumerator;
 
 pub struct PulsarSplitEnumerator {
     client: Pulsar<TokioExecutor>,