Skip to main content

risingwave_stream/from_proto/source/
fs_fetch.rs

1// Copyright 2023 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 risingwave_connector::WithOptionsSecResolved;
18use risingwave_connector::source::ConnectorProperties;
19use risingwave_connector::source::filesystem::opendal_source::{
20    OpendalAzblob, OpendalGcs, OpendalPosixFs, OpendalS3,
21};
22use risingwave_connector::source::reader::desc::SourceDescBuilder;
23use risingwave_pb::stream_plan::StreamFsFetchNode;
24use risingwave_storage::StateStore;
25
26use crate::error::StreamResult;
27use crate::executor::source::{
28    BatchAdbcSnowflakeFetchExecutor, BatchIcebergFetchExecutor, BatchOpendalFsFetchExecutor,
29    BatchPosixFsFetchExecutor, FsFetchExecutor, IcebergFetchExecutor, IcebergUpdateFetchExecutor,
30    SourceStateTableHandler, StreamSourceCore,
31};
32use crate::executor::{Execute, Executor};
33use crate::from_proto::ExecutorBuilder;
34use crate::from_proto::source::{iceberg_update_mode, is_full_reload_refresh};
35use crate::task::ExecutorParams;
36
37pub struct FsFetchExecutorBuilder;
38
39impl_stream_node_body!(StreamFsFetch(StreamFsFetchNode) => FsFetchExecutorBuilder);
40
41impl ExecutorBuilder for FsFetchExecutorBuilder {
42    type Node = StreamFsFetchNode;
43
44    async fn new_boxed_executor(
45        params: ExecutorParams,
46        node: &Self::Node,
47        store: impl StateStore,
48    ) -> StreamResult<Executor> {
49        let [upstream]: [_; 1] = params.input.try_into().unwrap();
50
51        let source = node.node_inner.as_ref().unwrap();
52        let is_full_reload_refresh = is_full_reload_refresh(&source.refresh_mode);
53
54        let source_id = source.source_id;
55        let source_name = source.source_name.clone();
56        let source_info = source.get_info()?;
57        let source_options_with_secret =
58            WithOptionsSecResolved::new(source.with_properties.clone(), source.secret_refs.clone());
59        let properties = ConnectorProperties::extract(source_options_with_secret.clone(), false)?;
60        let source_desc_builder = SourceDescBuilder::new(
61            source.columns.clone(),
62            params.env.source_metrics(),
63            source.row_id_index.map(|x| x as _),
64            source_options_with_secret,
65            source_info.clone(),
66            params.config.developer.connector_message_buffer_size,
67            params.info.stream_key.clone(),
68        );
69
70        let source_column_ids: Vec<_> = source_desc_builder
71            .column_catalogs_to_source_column_descs()
72            .iter()
73            .map(|column| column.column_id)
74            .collect();
75
76        let vnodes = Some(Arc::new(params.vnode_bitmap));
77        let is_iceberg_update = iceberg_update_mode(
78            &source.with_properties,
79            &source.refresh_mode,
80            source.rate_limit,
81            source.row_id_index,
82            params.info.stream_kind,
83        )?;
84        if is_iceberg_update {
85            crate::executor::source::validate_update_state_table(source.get_state_table()?, false)?;
86        }
87        let state_table_handler = SourceStateTableHandler::from_table_catalog_with_vnodes(
88            source.state_table.as_ref().unwrap(),
89            store.clone(),
90            vnodes,
91        )
92        .await;
93        let stream_source_core = StreamSourceCore::new(
94            source_id,
95            source_name,
96            source_column_ids,
97            source_desc_builder,
98            state_table_handler,
99        );
100
101        let exec = match properties {
102            risingwave_connector::source::ConnectorProperties::Gcs(_) => {
103                if is_full_reload_refresh {
104                    BatchOpendalFsFetchExecutor::<_, OpendalGcs>::new(
105                        params.actor_context.clone(),
106                        stream_source_core,
107                        upstream,
108                        source.rate_limit,
109                        params.local_barrier_manager.clone(),
110                        source.associated_table_id,
111                    )
112                    .boxed()
113                } else {
114                    FsFetchExecutor::<_, OpendalGcs>::new(
115                        params.actor_context.clone(),
116                        stream_source_core,
117                        upstream,
118                        source.rate_limit,
119                    )
120                    .boxed()
121                }
122            }
123            risingwave_connector::source::ConnectorProperties::OpendalS3(_) => {
124                if is_full_reload_refresh {
125                    BatchOpendalFsFetchExecutor::<_, OpendalS3>::new(
126                        params.actor_context.clone(),
127                        stream_source_core,
128                        upstream,
129                        source.rate_limit,
130                        params.local_barrier_manager.clone(),
131                        source.associated_table_id,
132                    )
133                    .boxed()
134                } else {
135                    FsFetchExecutor::<_, OpendalS3>::new(
136                        params.actor_context.clone(),
137                        stream_source_core,
138                        upstream,
139                        source.rate_limit,
140                    )
141                    .boxed()
142                }
143            }
144            risingwave_connector::source::ConnectorProperties::Iceberg(_) => {
145                if is_iceberg_update {
146                    let output_columns: Vec<risingwave_common::catalog::ColumnCatalog> =
147                        source.columns.iter().cloned().map(Into::into).collect();
148                    if output_columns
149                        .iter()
150                        .map(|column| column.data_type().clone())
151                        .collect::<Vec<_>>()
152                        != params.info.schema.data_types()
153                    {
154                        return Err(anyhow::anyhow!(
155                            "Iceberg update Fetch columns differ from its output schema"
156                        )
157                        .into());
158                    }
159                    IcebergUpdateFetchExecutor::new(
160                        params.actor_context.clone(),
161                        stream_source_core,
162                        upstream,
163                        params.config.clone(),
164                        output_columns,
165                        params.info.stream_key.clone(),
166                    )
167                    .boxed()
168                } else if is_full_reload_refresh {
169                    BatchIcebergFetchExecutor::new(
170                        params.actor_context.clone(),
171                        stream_source_core,
172                        upstream,
173                        params.local_barrier_manager.clone(),
174                        params.config.clone(),
175                        source.associated_table_id,
176                    )
177                    .boxed()
178                } else {
179                    IcebergFetchExecutor::new(
180                        params.actor_context.clone(),
181                        stream_source_core,
182                        upstream,
183                        source.rate_limit,
184                        params.config.clone(),
185                    )
186                    .boxed()
187                }
188            }
189            risingwave_connector::source::ConnectorProperties::Azblob(_) => {
190                FsFetchExecutor::<_, OpendalAzblob>::new(
191                    params.actor_context.clone(),
192                    stream_source_core,
193                    upstream,
194                    source.rate_limit,
195                )
196                .boxed()
197            }
198            risingwave_connector::source::ConnectorProperties::PosixFs(_) => {
199                FsFetchExecutor::<_, OpendalPosixFs>::new(
200                    params.actor_context.clone(),
201                    stream_source_core,
202                    upstream,
203                    source.rate_limit,
204                )
205                .boxed()
206            }
207            risingwave_connector::source::ConnectorProperties::BatchPosixFs(_) => {
208                BatchPosixFsFetchExecutor::new(
209                    params.actor_context.clone(),
210                    stream_source_core,
211                    upstream,
212                    source.rate_limit,
213                    params.local_barrier_manager.clone(),
214                    source.associated_table_id,
215                )
216                .boxed()
217            }
218            risingwave_connector::source::ConnectorProperties::AdbcSnowflake(_) => {
219                if is_full_reload_refresh {
220                    BatchAdbcSnowflakeFetchExecutor::new(
221                        params.actor_context.clone(),
222                        stream_source_core,
223                        upstream,
224                        params.local_barrier_manager.clone(),
225                        source.associated_table_id,
226                    )
227                    .boxed()
228                } else {
229                    unreachable!("AdbcSnowflake connector only supports FULL_RELOAD refresh mode")
230                }
231            }
232            _ => unreachable!(),
233        };
234        Ok((params.info, exec).into())
235    }
236}