risingwave_connector/source/reader/
reader.rs

1// Copyright 2024 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::sync::Arc;
16
17use anyhow::Context;
18use async_nats::jetstream::consumer::AckPolicy;
19use futures::StreamExt;
20use futures::stream::pending;
21use futures_async_stream::try_stream;
22use itertools::Itertools;
23use risingwave_common::bail;
24use risingwave_common::catalog::ColumnId;
25use thiserror_ext::AsReport as _;
26
27use crate::error::ConnectorResult;
28use crate::parser::{CommonParserConfig, ParserConfig, SpecificParserConfig};
29use crate::source::cdc::{CdcSourceRuntimeFlags, CdcSourceType};
30use crate::source::filesystem::opendal_source::opendal_enumerator::OpendalEnumerator;
31use crate::source::filesystem::opendal_source::{
32    DEFAULT_REFRESH_INTERVAL_SEC, OpendalAzblob, OpendalGcs, OpendalPosixFs, OpendalS3,
33    OpendalSource,
34};
35use crate::source::filesystem::{FsPageItem, OpendalFsSplit};
36use crate::source::{
37    BoxSourceReaderEventStream, BoxTryStream, Column, ConnectorProperties, ConnectorState,
38    CreateSplitReaderOpt, CreateSplitReaderResult, SourceColumnDesc, SourceContext,
39    WaitCheckpointTask,
40};
41use crate::{WithOptionsSecResolved, dispatch_cdc_source_type};
42
43#[derive(Clone, Debug)]
44pub struct SourceReader {
45    pub config: ConnectorProperties,
46    pub columns: Vec<SourceColumnDesc>,
47    pub parser_config: SpecificParserConfig,
48    pub connector_message_buffer_size: usize,
49}
50
51impl SourceReader {
52    pub fn new(
53        properties: WithOptionsSecResolved,
54        columns: Vec<SourceColumnDesc>,
55        connector_message_buffer_size: usize,
56        parser_config: SpecificParserConfig,
57    ) -> ConnectorResult<Self> {
58        let config = ConnectorProperties::extract(properties, false)?;
59
60        Ok(Self {
61            config,
62            columns,
63            parser_config,
64            connector_message_buffer_size,
65        })
66    }
67
68    fn get_target_columns(
69        &self,
70        column_ids: Vec<ColumnId>,
71    ) -> ConnectorResult<Vec<SourceColumnDesc>> {
72        column_ids
73            .iter()
74            .map(|id| {
75                self.columns
76                    .iter()
77                    .find(|c| c.column_id == *id)
78                    .with_context(|| {
79                        format!("Failed to find column id: {} in source: {:?}", id, self)
80                    })
81                    .cloned()
82            })
83            .try_collect()
84            .map_err(Into::into)
85    }
86
87    pub fn get_source_list(&self) -> ConnectorResult<BoxTryStream<FsPageItem>> {
88        let config = self.config.clone();
89        let list_interval_sec: u64;
90        let get_list_interval_sec =
91            |interval: Option<u64>| -> u64 { interval.unwrap_or(DEFAULT_REFRESH_INTERVAL_SEC) };
92        match config {
93            ConnectorProperties::Gcs(prop) => {
94                list_interval_sec = get_list_interval_sec(prop.fs_common.refresh_interval_sec);
95                let lister: OpendalEnumerator<OpendalGcs> =
96                    OpendalEnumerator::new_gcs_source(*prop)?;
97                Ok(build_opendal_fs_list_stream(lister, list_interval_sec))
98            }
99            ConnectorProperties::OpendalS3(prop) => {
100                list_interval_sec = get_list_interval_sec(prop.fs_common.refresh_interval_sec);
101                let lister: OpendalEnumerator<OpendalS3> = OpendalEnumerator::new_s3_source(
102                    &prop.s3_properties,
103                    prop.assume_role,
104                    prop.fs_common.compression_format,
105                )?;
106                Ok(build_opendal_fs_list_stream(lister, list_interval_sec))
107            }
108            ConnectorProperties::Azblob(prop) => {
109                list_interval_sec = get_list_interval_sec(prop.fs_common.refresh_interval_sec);
110                let lister: OpendalEnumerator<OpendalAzblob> =
111                    OpendalEnumerator::new_azblob_source(*prop)?;
112                Ok(build_opendal_fs_list_stream(lister, list_interval_sec))
113            }
114            ConnectorProperties::PosixFs(prop) => {
115                list_interval_sec = get_list_interval_sec(prop.fs_common.refresh_interval_sec);
116                let lister: OpendalEnumerator<OpendalPosixFs> =
117                    OpendalEnumerator::new_posix_fs_source(*prop)?;
118                Ok(build_opendal_fs_list_stream(lister, list_interval_sec))
119            }
120            other => bail!("Unsupported source: {:?}", other),
121        }
122    }
123
124    /// Refer to `WaitCheckpointWorker` for more details.
125    pub async fn create_wait_checkpoint_task(&self) -> ConnectorResult<Option<WaitCheckpointTask>> {
126        if dispatch_cdc_source_type!(
127            CdcSourceType::from(&self.config),
128            T,
129            T::NEEDS_OFFSET_COMMIT_ON_CHECKPOINT,
130            false
131        ) {
132            return Ok(Some(WaitCheckpointTask::CommitCdcOffset(None)));
133        }
134        Ok(match &self.config {
135            ConnectorProperties::GooglePubsub(prop) => Some(WaitCheckpointTask::AckPubsubMessage(
136                prop.subscription_client().await?,
137                vec![],
138            )),
139            ConnectorProperties::Nats(prop) => {
140                match prop.nats_properties_consumer.get_ack_policy()? {
141                    a @ AckPolicy::Explicit | a @ AckPolicy::All => {
142                        Some(WaitCheckpointTask::AckNatsJetStream(
143                            prop.common.build_context().await?,
144                            vec![],
145                            a,
146                        ))
147                    }
148                    AckPolicy::None | AckPolicy::FlowControl => None,
149                }
150            }
151            ConnectorProperties::Pulsar(_) => Some(WaitCheckpointTask::AckPulsarMessage(vec![])),
152            _ => None,
153        })
154    }
155
156    /// Build `SplitReader`s and then `BoxSourceReaderEventStream` from the given `ConnectorState` (`SplitImpl`s).
157    ///
158    /// If `seek_to_latest` is true, will also return the latest splits after seek.
159    pub async fn build_stream(
160        &self,
161        state: ConnectorState,
162        column_ids: Vec<ColumnId>,
163        source_ctx: Arc<SourceContext>,
164        seek_to_latest: bool,
165    ) -> ConnectorResult<(BoxSourceReaderEventStream, CreateSplitReaderResult)> {
166        let Some(splits) = state else {
167            return Ok((pending().boxed(), Default::default()));
168        };
169        let config = self.config.clone();
170        let columns = self.get_target_columns(column_ids)?;
171
172        let data_gen_columns = Some(
173            columns
174                .iter()
175                .map(|col| Column {
176                    name: col.name.clone(),
177                    data_type: col.data_type.clone(),
178                    is_visible: col.is_visible(),
179                })
180                .collect_vec(),
181        );
182
183        let parser_config = ParserConfig {
184            specific: self.parser_config.clone(),
185            common: CommonParserConfig {
186                rw_columns: columns,
187            },
188        };
189
190        config
191            .create_split_reader(
192                splits,
193                parser_config,
194                source_ctx,
195                data_gen_columns,
196                CreateSplitReaderOpt {
197                    seek_to_latest,
198                    ..Default::default()
199                },
200            )
201            .await
202    }
203}
204
205#[try_stream(boxed, ok = FsPageItem, error = crate::error::ConnectorError)]
206async fn build_opendal_fs_list_stream<Src: OpendalSource>(
207    lister: OpendalEnumerator<Src>,
208    list_interval_sec: u64,
209) {
210    loop {
211        let matcher = lister.get_matcher();
212        let mut object_metadata_iter = lister.list().await?;
213        while let Some(list_res) = object_metadata_iter.next().await {
214            match list_res {
215                Ok(res) => {
216                    if matcher
217                        .as_ref()
218                        .map(|m| m.matches(&res.name) || m.to_string() == res.name)
219                        .unwrap_or(true)
220                    {
221                        yield res
222                    } else {
223                        continue;
224                    }
225                }
226                Err(err) => {
227                    tracing::error!(error = %err.as_report(), "list object fail");
228                    return Err(err);
229                }
230            }
231        }
232        tokio::time::sleep(std::time::Duration::from_secs(list_interval_sec)).await;
233    }
234}
235
236#[try_stream(boxed, ok = OpendalFsSplit<Src>,  error = crate::error::ConnectorError)]
237pub async fn build_opendal_fs_list_for_batch<Src: OpendalSource>(lister: OpendalEnumerator<Src>) {
238    let matcher = lister.get_matcher();
239    let mut object_metadata_iter = lister.list().await?;
240
241    while let Some(list_res) = object_metadata_iter.next().await {
242        match list_res {
243            Ok(res) => {
244                if matcher
245                    .as_ref()
246                    .map(|m| m.matches(&res.name))
247                    .unwrap_or(true)
248                {
249                    let split = OpendalFsSplit::new(res.name, 0, res.size as usize);
250                    yield split
251                } else {
252                    continue;
253                }
254            }
255            Err(err) => {
256                tracing::error!(error = %err.as_report(), "list object fail");
257                return Err(err);
258            }
259        }
260    }
261}
risingwave_connector/source/reader/reader.rs

risingwave_connector/source/reader/
reader.rs
