Skip to main content

risingwave_connector/
with_options.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::collections::{BTreeMap, BTreeSet, HashMap};
16use std::marker::PhantomData;
17use std::time::Duration;
18
19use risingwave_pb::id::SecretId;
20use risingwave_pb::secret::PbSecretRef;
21
22use crate::dispatch_cdc_source_type;
23use crate::error::ConnectorResult;
24use crate::sink::catalog::SinkFormatDesc;
25use crate::source::cdc::{CdcSourceRuntimeFlags, CdcSourceType};
26use crate::source::iceberg::ICEBERG_CONNECTOR;
27use crate::source::{
28    ADBC_SNOWFLAKE_CONNECTOR, AZBLOB_CONNECTOR, BATCH_POSIX_FS_CONNECTOR, GCS_CONNECTOR,
29    KAFKA_CONNECTOR, LEGACY_S3_CONNECTOR, OPENDAL_S3_CONNECTOR, POSIX_FS_CONNECTOR,
30    PULSAR_CONNECTOR, UPSTREAM_SOURCE_KEY,
31};
32
33/// Marker trait for `WITH` options. Only for `#[derive(WithOptions)]`, should not be used manually.
34///
35/// This is used to ensure the `WITH` options types have reasonable structure.
36///
37/// TODO: add this bound for sink. There's a `SourceProperties` trait for sources, but no similar
38/// things for sinks.
39pub trait WithOptions {
40    #[doc(hidden)]
41    #[inline(always)]
42    fn assert_receiver_is_with_options(&self) {}
43}
44
45// Currently CDC properties are handled specially.
46// - It simply passes HashMap to Java DBZ.
47// - It's not handled by serde.
48// - It contains fields other than WITH options.
49// TODO: remove the workaround here. And also use #[derive] for it.
50
51impl<T: crate::source::cdc::CdcSourceTypeTrait> WithOptions
52    for crate::source::cdc::CdcProperties<T>
53{
54}
55
56// impl the trait for value types
57
58impl<T: WithOptions> WithOptions for Option<T> {}
59impl WithOptions for Vec<String> {}
60impl WithOptions for Vec<u64> {}
61impl WithOptions for HashMap<String, String> {}
62impl WithOptions for BTreeMap<String, String> {}
63
64impl WithOptions for String {}
65impl WithOptions for bool {}
66impl WithOptions for usize {}
67impl WithOptions for u8 {}
68impl WithOptions for u16 {}
69impl WithOptions for u32 {}
70impl WithOptions for u64 {}
71impl WithOptions for i32 {}
72impl WithOptions for i64 {}
73impl WithOptions for f64 {}
74impl WithOptions for std::time::Duration {}
75impl WithOptions for crate::connector_common::MqttQualityOfService {}
76impl WithOptions for crate::connector_common::SslMode {}
77impl WithOptions for crate::sink::file_sink::opendal_sink::PathPartitionPrefix {}
78impl WithOptions for crate::sink::kafka::CompressionCodec {}
79impl WithOptions for crate::sink::pulsar::PulsarRoutingMode {}
80impl WithOptions for crate::source::filesystem::file_common::CompressionFormat {}
81impl WithOptions for nexmark::config::RateShape {}
82impl WithOptions for nexmark::event::EventType {}
83impl<T> WithOptions for PhantomData<T> {}
84
85pub trait Get {
86    fn get(&self, key: &str) -> Option<&String>;
87}
88
89pub trait GetKeyIter {
90    fn key_iter(&self) -> impl Iterator<Item = &str>;
91}
92
93impl GetKeyIter for HashMap<String, String> {
94    fn key_iter(&self) -> impl Iterator<Item = &str> {
95        self.keys().map(|s| s.as_str())
96    }
97}
98
99impl Get for HashMap<String, String> {
100    fn get(&self, key: &str) -> Option<&String> {
101        self.get(key)
102    }
103}
104
105impl Get for BTreeMap<String, String> {
106    fn get(&self, key: &str) -> Option<&String> {
107        self.get(key)
108    }
109}
110
111impl GetKeyIter for BTreeMap<String, String> {
112    fn key_iter(&self) -> impl Iterator<Item = &str> {
113        self.keys().map(|s| s.as_str())
114    }
115}
116
117/// Utility methods for `WITH` properties (`HashMap` and `BTreeMap`).
118pub trait WithPropertiesExt: Get + GetKeyIter + Sized {
119    #[inline(always)]
120    fn get_connector(&self) -> Option<String> {
121        self.get(UPSTREAM_SOURCE_KEY).map(|s| s.to_lowercase())
122    }
123
124    #[inline(always)]
125    fn is_kafka_connector(&self) -> bool {
126        let Some(connector) = self.get_connector() else {
127            return false;
128        };
129        connector == KAFKA_CONNECTOR
130    }
131
132    #[inline(always)]
133    fn is_pulsar_connector(&self) -> bool {
134        let Some(connector) = self.get_connector() else {
135            return false;
136        };
137        connector == PULSAR_CONNECTOR
138    }
139
140    #[inline(always)]
141    fn get_sync_call_timeout(&self) -> Option<Duration> {
142        const SYNC_CALL_TIMEOUT_KEY: &str = "properties.sync.call.timeout"; // only from kafka props, add more if needed
143        self.get(SYNC_CALL_TIMEOUT_KEY)
144            // ignore the error is ok here, because we will parse the field again when building the properties and has more precise error message
145            .and_then(|s| duration_str::parse_std(s).ok())
146    }
147
148    #[inline(always)]
149    fn is_cdc_connector(&self) -> bool {
150        let Some(connector) = self.get_connector() else {
151            return false;
152        };
153        connector.contains("-cdc")
154    }
155
156    fn cdc_source_type(&self) -> CdcSourceType {
157        CdcSourceType::from(self.get_connector().unwrap_or_default().as_str())
158    }
159
160    /// It is shared when `CREATE SOURCE`, and not shared when `CREATE TABLE`. So called "shareable".
161    fn is_shareable_cdc_connector(&self) -> bool {
162        dispatch_cdc_source_type!(self.cdc_source_type(), T, T::CAN_BACKFILL, false)
163    }
164
165    fn enable_transaction_metadata(&self) -> bool {
166        dispatch_cdc_source_type!(
167            self.cdc_source_type(),
168            T,
169            T::CAN_USE_TRANSACTION_METADATA,
170            false
171        )
172    }
173
174    fn is_shareable_non_cdc_connector(&self) -> bool {
175        self.is_kafka_connector()
176    }
177
178    #[inline(always)]
179    fn is_iceberg_connector(&self) -> bool {
180        let Some(connector) = self.get_connector() else {
181            return false;
182        };
183        connector == ICEBERG_CONNECTOR
184    }
185
186    fn connector_need_pk(&self) -> bool {
187        // Legacy Iceberg scans do not require a key. Retracting update ingestion uses
188        // the SQL primary key, like other non-append-only sources.
189        !self.is_iceberg_connector() || self.is_iceberg_update_source()
190    }
191
192    fn is_iceberg_update_source(&self) -> bool {
193        self.is_iceberg_connector()
194            && self
195                .get(crate::source::iceberg::STREAMING_UPDATES_KEY)
196                .is_some_and(|value| value.eq_ignore_ascii_case("true"))
197    }
198
199    fn is_legacy_fs_connector(&self) -> bool {
200        self.get(UPSTREAM_SOURCE_KEY)
201            .map(|s| s.eq_ignore_ascii_case(LEGACY_S3_CONNECTOR))
202            .unwrap_or(false)
203    }
204
205    fn is_new_fs_connector(&self) -> bool {
206        self.get(UPSTREAM_SOURCE_KEY)
207            .map(|s| {
208                s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR)
209                    || s.eq_ignore_ascii_case(POSIX_FS_CONNECTOR)
210                    || s.eq_ignore_ascii_case(GCS_CONNECTOR)
211                    || s.eq_ignore_ascii_case(AZBLOB_CONNECTOR)
212            })
213            .unwrap_or(false)
214    }
215
216    fn is_batch_connector(&self) -> bool {
217        self.get(UPSTREAM_SOURCE_KEY)
218            .map(|s| {
219                s.eq_ignore_ascii_case(BATCH_POSIX_FS_CONNECTOR)
220                    || s.eq_ignore_ascii_case(ADBC_SNOWFLAKE_CONNECTOR)
221            })
222            .unwrap_or(false)
223    }
224
225    fn supports_full_reload_refresh(&self) -> bool {
226        self.get(UPSTREAM_SOURCE_KEY)
227            .map(|s| {
228                s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR)
229                    || s.eq_ignore_ascii_case(GCS_CONNECTOR)
230                    || s.eq_ignore_ascii_case(BATCH_POSIX_FS_CONNECTOR)
231                    || s.eq_ignore_ascii_case(ICEBERG_CONNECTOR)
232                    || s.eq_ignore_ascii_case(ADBC_SNOWFLAKE_CONNECTOR)
233            })
234            .unwrap_or(false)
235    }
236
237    fn requires_singleton(&self) -> bool {
238        self.is_new_fs_connector() || self.is_iceberg_connector() || self.is_batch_connector()
239    }
240}
241
242impl<T: Get + GetKeyIter> WithPropertiesExt for T {}
243
244/// Options or properties extracted from the `WITH` clause of DDLs.
245#[derive(Default, Clone, Debug, PartialEq, Eq, Hash)]
246pub struct WithOptionsSecResolved {
247    inner: BTreeMap<String, String>,
248    secret_ref: BTreeMap<String, PbSecretRef>,
249}
250
251impl std::ops::Deref for WithOptionsSecResolved {
252    type Target = BTreeMap<String, String>;
253
254    fn deref(&self) -> &Self::Target {
255        &self.inner
256    }
257}
258
259impl std::ops::DerefMut for WithOptionsSecResolved {
260    fn deref_mut(&mut self) -> &mut Self::Target {
261        &mut self.inner
262    }
263}
264
265impl WithOptionsSecResolved {
266    /// Create a new [`WithOptions`] from a option [`BTreeMap`] and resolved secret ref.
267    pub fn new(inner: BTreeMap<String, String>, secret_ref: BTreeMap<String, PbSecretRef>) -> Self {
268        Self { inner, secret_ref }
269    }
270
271    pub fn as_plaintext(&self) -> &BTreeMap<String, String> {
272        &self.inner
273    }
274
275    pub fn as_secret(&self) -> &BTreeMap<String, PbSecretRef> {
276        &self.secret_ref
277    }
278
279    pub fn handle_update(
280        &mut self,
281        update_alter_props: BTreeMap<String, String>,
282        update_alter_secret_refs: BTreeMap<String, PbSecretRef>,
283    ) -> ConnectorResult<(Vec<SecretId>, Vec<SecretId>)> {
284        let old_secret_ids = self
285            .secret_ref
286            .values()
287            .map(|secret_ref| secret_ref.secret_id)
288            .collect::<BTreeSet<_>>();
289
290        // make sure the key in update_alter_props and update_alter_secret_refs not collide
291        for key in update_alter_props.keys() {
292            if update_alter_secret_refs.contains_key(key) {
293                return Err(
294                    anyhow::anyhow!("the key {} is set both in plaintext and secret", key).into(),
295                );
296            }
297        }
298
299        // remove legacy key if it's set in both plaintext and secret
300        // When a property changes from secret to plaintext, remove the old secret dependency
301        for k in update_alter_props.keys() {
302            self.secret_ref.remove(k);
303        }
304
305        // Handle secret ref updates
306        for k in update_alter_secret_refs.keys() {
307            // Remove any plaintext value for this key
308            self.inner.remove(k);
309        }
310
311        self.inner.extend(update_alter_props);
312        self.secret_ref.extend(update_alter_secret_refs);
313
314        let new_secret_ids = self
315            .secret_ref
316            .values()
317            .map(|secret_ref| secret_ref.secret_id)
318            .collect::<BTreeSet<_>>();
319        let to_add_secret_dep = new_secret_ids
320            .difference(&old_secret_ids)
321            .copied()
322            .collect();
323        let to_remove_secret_dep = old_secret_ids
324            .difference(&new_secret_ids)
325            .copied()
326            .collect();
327
328        Ok((to_add_secret_dep, to_remove_secret_dep))
329    }
330
331    /// Create a new [`WithOptions`] from a [`BTreeMap`].
332    pub fn without_secrets(inner: BTreeMap<String, String>) -> Self {
333        Self {
334            inner,
335            secret_ref: Default::default(),
336        }
337    }
338
339    /// Take the value of the option map and secret refs.
340    pub fn into_parts(self) -> (BTreeMap<String, String>, BTreeMap<String, PbSecretRef>) {
341        (self.inner, self.secret_ref)
342    }
343
344    pub fn value_eq_ignore_case(&self, key: &str, val: &str) -> bool {
345        if let Some(inner_val) = self.inner.get(key)
346            && inner_val.eq_ignore_ascii_case(val)
347        {
348            return true;
349        }
350        false
351    }
352}
353
354/// For `planner_test` crate so that it does not depend directly on `connector` crate just for `SinkFormatDesc`.
355impl TryFrom<&WithOptionsSecResolved> for Option<SinkFormatDesc> {
356    type Error = crate::sink::SinkError;
357
358    fn try_from(value: &WithOptionsSecResolved) -> std::result::Result<Self, Self::Error> {
359        let connector = value.get(crate::sink::CONNECTOR_TYPE_KEY);
360        let r#type = value.get(crate::sink::SINK_TYPE_OPTION);
361        match (connector, r#type) {
362            (Some(c), Some(t)) => SinkFormatDesc::from_legacy_type(c, t),
363            _ => Ok(None),
364        }
365    }
366}
367
368impl Get for WithOptionsSecResolved {
369    fn get(&self, key: &str) -> Option<&String> {
370        self.inner.get(key)
371    }
372}
373
374impl GetKeyIter for WithOptionsSecResolved {
375    fn key_iter(&self) -> impl Iterator<Item = &str> {
376        self.inner.keys().map(|s| s.as_str())
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    fn source_options(connector: &str) -> WithOptionsSecResolved {
385        WithOptionsSecResolved::without_secrets(BTreeMap::from([(
386            UPSTREAM_SOURCE_KEY.to_owned(),
387            connector.to_owned(),
388        )]))
389    }
390
391    #[test]
392    fn test_full_reload_refresh_connector_whitelist() {
393        for connector in [
394            OPENDAL_S3_CONNECTOR,
395            GCS_CONNECTOR,
396            BATCH_POSIX_FS_CONNECTOR,
397            ICEBERG_CONNECTOR,
398            ADBC_SNOWFLAKE_CONNECTOR,
399        ] {
400            assert!(
401                source_options(connector).supports_full_reload_refresh(),
402                "{connector} should support FULL_RELOAD refresh"
403            );
404        }
405
406        for connector in [KAFKA_CONNECTOR, POSIX_FS_CONNECTOR, AZBLOB_CONNECTOR] {
407            assert!(
408                !source_options(connector).supports_full_reload_refresh(),
409                "{connector} should not support FULL_RELOAD refresh"
410            );
411        }
412    }
413}