refactor(source): refine some code of split reader

src/connector/src/macros.rs

@@ -235,7 +235,7 @@ macro_rules! impl_common_split_reader_logic {
     ($reader:ty, $props:ty) => {
         impl $reader {
             #[try_stream(boxed, ok = $crate::source::StreamChunkWithState, error = risingwave_common::error::RwError)]
-            pub(crate) async fn into_msg_stream(self) {
+            pub(crate) async fn into_chunk_stream(self) {
                 let parser_config = self.parser_config.clone();
                 let actor_id = self.source_info.actor_id.to_string();
                 let source_id = self.source_info.source_id.to_string();
@@ -272,33 +272,5 @@ macro_rules! impl_common_split_reader_logic {
                 }
             }
         }
-
-        #[async_trait]
-        impl $crate::source::SplitReaderV2 for $reader {
-            type Properties = $props;
-
-            async fn new(
-                properties: $props,
-                state: Vec<SplitImpl>,
-                parser_config: ParserConfig,
-                metrics: Arc<SourceMetrics>,
-                source_info: SourceInfo,
-                columns: Option<Vec<Column>>,
-            ) -> Result<Self> {
-                Self::new(
-                    properties,
-                    state,
-                    parser_config,
-                    metrics,
-                    source_info,
-                    columns,
-                )
-                .await
-            }
-
-            fn into_stream(self) -> $crate::source::BoxSourceWithStateStream {
-                self.into_msg_stream()
-            }
-        }
     };
 }

src/connector/src/source/cdc/source/reader.rs

@@ -28,7 +28,9 @@ use crate::parser::ParserConfig;
 use crate::source::base::SourceMessage;
 use crate::source::cdc::CdcProperties;
 use crate::source::monitor::SourceMetrics;
-use crate::source::{Column, SourceInfo, SplitId, SplitImpl, SplitMetaData};
+use crate::source::{
+    BoxSourceWithStateStream, Column, SourceInfo, SplitId, SplitImpl, SplitMetaData, SplitReaderV2,
+};
 
 impl_common_split_reader_logic!(CdcSplitReader, CdcProperties);
 
@@ -43,7 +45,10 @@ pub struct CdcSplitReader {
     source_info: SourceInfo,
 }
 
-impl CdcSplitReader {
+#[async_trait]
+impl SplitReaderV2 for CdcSplitReader {
+    type Properties = CdcProperties;
+
     #[allow(clippy::unused_async)]
     async fn new(
         conn_props: CdcProperties,
@@ -71,6 +76,10 @@ impl CdcSplitReader {
             )),
         }
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }
 
 impl CdcSplitReader {

src/connector/src/source/datagen/source/reader.rs

@@ -30,7 +30,8 @@ use crate::source::datagen::source::SEQUENCE_FIELD_KIND;
 use crate::source::datagen::{DatagenProperties, DatagenSplit};
 use crate::source::monitor::SourceMetrics;
 use crate::source::{
-    BoxSourceStream, Column, DataType, SourceInfo, SplitId, SplitImpl, SplitMetaData,
+    BoxSourceStream, BoxSourceWithStateStream, Column, DataType, SourceInfo, SplitId, SplitImpl,
+    SplitMetaData, SplitReaderV2,
 };
 
 impl_common_split_reader_logic!(DatagenSplitReader, DatagenProperties);
@@ -45,7 +46,10 @@ pub struct DatagenSplitReader {
     source_info: SourceInfo,
 }
 
-impl DatagenSplitReader {
+#[async_trait]
+impl SplitReaderV2 for DatagenSplitReader {
+    type Properties = DatagenProperties;
+
     #[allow(clippy::unused_async)]
     async fn new(
         properties: DatagenProperties,
@@ -132,6 +136,12 @@ impl DatagenSplitReader {
         })
     }
 
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
+}
+
+impl DatagenSplitReader {
     pub(crate) fn into_data_stream(self) -> BoxSourceStream {
         // Will buffer at most 4 event chunks.
         const BUFFER_SIZE: usize = 4;

src/connector/src/source/google_pubsub/source/reader.rs

@@ -29,7 +29,10 @@ use crate::impl_common_split_reader_logic;
 use crate::parser::ParserConfig;
 use crate::source::google_pubsub::PubsubProperties;
 use crate::source::monitor::SourceMetrics;
-use crate::source::{Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData};
+use crate::source::{
+    BoxSourceWithStateStream, Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData,
+    SplitReaderV2,
+};
 
 const PUBSUB_MAX_FETCH_MESSAGES: usize = 1024;
 
@@ -109,7 +112,10 @@ impl PubsubSplitReader {
     }
 }
 
-impl PubsubSplitReader {
+#[async_trait]
+impl SplitReaderV2 for PubsubSplitReader {
+    type Properties = PubsubProperties;
+
     async fn new(
         properties: PubsubProperties,
         splits: Vec<SplitImpl>,
@@ -169,4 +175,8 @@ impl PubsubSplitReader {
             source_info,
         })
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }

src/connector/src/source/kafka/source/reader.rs

@@ -29,7 +29,9 @@ use crate::parser::ParserConfig;
 use crate::source::base::{SourceMessage, MAX_CHUNK_SIZE};
 use crate::source::kafka::KafkaProperties;
 use crate::source::monitor::SourceMetrics;
-use crate::source::{Column, SourceInfo, SplitId, SplitImpl, SplitMetaData};
+use crate::source::{
+    BoxSourceWithStateStream, Column, SourceInfo, SplitId, SplitImpl, SplitMetaData, SplitReaderV2,
+};
 
 impl_common_split_reader_logic!(KafkaSplitReader, KafkaProperties);
 
@@ -46,7 +48,10 @@ pub struct KafkaSplitReader {
     source_info: SourceInfo,
 }
 
-impl KafkaSplitReader {
+#[async_trait]
+impl SplitReaderV2 for KafkaSplitReader {
+    type Properties = KafkaProperties;
+
     async fn new(
         properties: KafkaProperties,
         splits: Vec<SplitImpl>,
@@ -136,6 +141,10 @@ impl KafkaSplitReader {
             source_info,
         })
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }
 
 impl KafkaSplitReader {

src/connector/src/source/kinesis/source/reader.rs

@@ -32,7 +32,10 @@ use crate::source::kinesis::source::message::KinesisMessage;
 use crate::source::kinesis::split::KinesisOffset;
 use crate::source::kinesis::KinesisProperties;
 use crate::source::monitor::SourceMetrics;
-use crate::source::{Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData};
+use crate::source::{
+    BoxSourceWithStateStream, Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData,
+    SplitReaderV2,
+};
 
 impl_common_split_reader_logic!(KinesisSplitReader, KinesisProperties);
 
@@ -52,7 +55,10 @@ pub struct KinesisSplitReader {
     source_info: SourceInfo,
 }
 
-impl KinesisSplitReader {
+#[async_trait]
+impl SplitReaderV2 for KinesisSplitReader {
+    type Properties = KinesisProperties;
+
     async fn new(
         properties: KinesisProperties,
         splits: Vec<SplitImpl>,
@@ -106,6 +112,10 @@ impl KinesisSplitReader {
             source_info,
         })
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }
 
 impl KinesisSplitReader {

src/connector/src/source/nexmark/source/reader.rs

@@ -31,7 +31,8 @@ use crate::source::monitor::SourceMetrics;
 use crate::source::nexmark::source::message::NexmarkMessage;
 use crate::source::nexmark::{NexmarkProperties, NexmarkSplit};
 use crate::source::{
-    BoxSourceStream, Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData,
+    BoxSourceStream, BoxSourceWithStateStream, Column, SourceInfo, SourceMessage, SplitId,
+    SplitImpl, SplitMetaData, SplitReaderV2,
 };
 
 impl_common_split_reader_logic!(NexmarkSplitReader, NexmarkProperties);
@@ -52,7 +53,10 @@ pub struct NexmarkSplitReader {
     source_info: SourceInfo,
 }
 
-impl NexmarkSplitReader {
+#[async_trait]
+impl SplitReaderV2 for NexmarkSplitReader {
+    type Properties = NexmarkProperties;
+
     #[allow(clippy::unused_async)]
     async fn new(
         properties: NexmarkProperties,
@@ -96,6 +100,10 @@ impl NexmarkSplitReader {
             source_info,
         })
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }
 
 impl NexmarkSplitReader {

src/connector/src/source/pulsar/source/reader.rs

@@ -31,7 +31,8 @@ use crate::source::monitor::SourceMetrics;
 use crate::source::pulsar::split::PulsarSplit;
 use crate::source::pulsar::{PulsarEnumeratorOffset, PulsarProperties};
 use crate::source::{
-    Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData, MAX_CHUNK_SIZE,
+    BoxSourceWithStateStream, Column, SourceInfo, SourceMessage, SplitId, SplitImpl, SplitMetaData,
+    SplitReaderV2, MAX_CHUNK_SIZE,
 };
 
 impl_common_split_reader_logic!(PulsarSplitReader, PulsarProperties);
@@ -89,7 +90,10 @@ fn parse_message_id(id: &str) -> Result<MessageIdData> {
     Ok(message_id)
 }
 
-impl PulsarSplitReader {
+#[async_trait]
+impl SplitReaderV2 for PulsarSplitReader {
+    type Properties = PulsarProperties;
+
     async fn new(
         props: PulsarProperties,
         splits: Vec<SplitImpl>,
@@ -157,6 +161,10 @@ impl PulsarSplitReader {
             source_info,
         })
     }
+
+    fn into_stream(self) -> BoxSourceWithStateStream {
+        self.into_chunk_stream()
+    }
 }
 
 impl PulsarSplitReader {