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::error::ConnectorResult;
23use crate::sink::catalog::SinkFormatDesc;
24use crate::source::cdc::MYSQL_CDC_CONNECTOR;
25use crate::source::cdc::external::ExternalCdcTableType;
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::sink::file_sink::opendal_sink::PathPartitionPrefix {}
77impl WithOptions for crate::sink::kafka::CompressionCodec {}
78impl WithOptions for crate::sink::pulsar::PulsarRoutingMode {}
79impl WithOptions for crate::source::filesystem::file_common::CompressionFormat {}
80impl WithOptions for nexmark::config::RateShape {}
81impl WithOptions for nexmark::event::EventType {}
82impl<T> WithOptions for PhantomData<T> {}
83
84pub trait Get {
85    fn get(&self, key: &str) -> Option<&String>;
86}
87
88pub trait GetKeyIter {
89    fn key_iter(&self) -> impl Iterator<Item = &str>;
90}
91
92impl GetKeyIter for HashMap<String, String> {
93    fn key_iter(&self) -> impl Iterator<Item = &str> {
94        self.keys().map(|s| s.as_str())
95    }
96}
97
98impl Get for HashMap<String, String> {
99    fn get(&self, key: &str) -> Option<&String> {
100        self.get(key)
101    }
102}
103
104impl Get for BTreeMap<String, String> {
105    fn get(&self, key: &str) -> Option<&String> {
106        self.get(key)
107    }
108}
109
110impl GetKeyIter for BTreeMap<String, String> {
111    fn key_iter(&self) -> impl Iterator<Item = &str> {
112        self.keys().map(|s| s.as_str())
113    }
114}
115
116/// Utility methods for `WITH` properties (`HashMap` and `BTreeMap`).
117pub trait WithPropertiesExt: Get + GetKeyIter + Sized {
118    #[inline(always)]
119    fn get_connector(&self) -> Option<String> {
120        self.get(UPSTREAM_SOURCE_KEY).map(|s| s.to_lowercase())
121    }
122
123    #[inline(always)]
124    fn is_kafka_connector(&self) -> bool {
125        let Some(connector) = self.get_connector() else {
126            return false;
127        };
128        connector == KAFKA_CONNECTOR
129    }
130
131    #[inline(always)]
132    fn is_pulsar_connector(&self) -> bool {
133        let Some(connector) = self.get_connector() else {
134            return false;
135        };
136        connector == PULSAR_CONNECTOR
137    }
138
139    #[inline(always)]
140    fn is_mysql_cdc_connector(&self) -> bool {
141        let Some(connector) = self.get_connector() else {
142            return false;
143        };
144        connector == MYSQL_CDC_CONNECTOR
145    }
146
147    #[inline(always)]
148    fn get_sync_call_timeout(&self) -> Option<Duration> {
149        const SYNC_CALL_TIMEOUT_KEY: &str = "properties.sync.call.timeout"; // only from kafka props, add more if needed
150        self.get(SYNC_CALL_TIMEOUT_KEY)
151            // ignore the error is ok here, because we will parse the field again when building the properties and has more precise error message
152            .and_then(|s| duration_str::parse_std(s).ok())
153    }
154
155    #[inline(always)]
156    fn is_cdc_connector(&self) -> bool {
157        let Some(connector) = self.get_connector() else {
158            return false;
159        };
160        connector.contains("-cdc")
161    }
162
163    /// It is shared when `CREATE SOURCE`, and not shared when `CREATE TABLE`. So called "shareable".
164    fn is_shareable_cdc_connector(&self) -> bool {
165        self.is_cdc_connector() && ExternalCdcTableType::from_properties(self).can_backfill()
166    }
167
168    fn enable_transaction_metadata(&self) -> bool {
169        ExternalCdcTableType::from_properties(self).enable_transaction_metadata()
170    }
171
172    fn is_shareable_non_cdc_connector(&self) -> bool {
173        self.is_kafka_connector()
174    }
175
176    #[inline(always)]
177    fn is_iceberg_connector(&self) -> bool {
178        let Some(connector) = self.get_connector() else {
179            return false;
180        };
181        connector == ICEBERG_CONNECTOR
182    }
183
184    fn connector_need_pk(&self) -> bool {
185        // Currently only iceberg connector doesn't need primary key
186        // introduced in https://github.com/risingwavelabs/risingwave/pull/14971
187        // XXX: This seems not the correct way. Iceberg doesn't necessarily lack a PK.
188        // "batch source" doesn't need a PK?
189        // For streaming, if it has a PK, do we want to use it? It seems not safe.
190        !self.is_iceberg_connector()
191    }
192
193    fn is_legacy_fs_connector(&self) -> bool {
194        self.get(UPSTREAM_SOURCE_KEY)
195            .map(|s| s.eq_ignore_ascii_case(LEGACY_S3_CONNECTOR))
196            .unwrap_or(false)
197    }
198
199    fn is_new_fs_connector(&self) -> bool {
200        self.get(UPSTREAM_SOURCE_KEY)
201            .map(|s| {
202                s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR)
203                    || s.eq_ignore_ascii_case(POSIX_FS_CONNECTOR)
204                    || s.eq_ignore_ascii_case(GCS_CONNECTOR)
205                    || s.eq_ignore_ascii_case(AZBLOB_CONNECTOR)
206            })
207            .unwrap_or(false)
208    }
209
210    fn is_batch_connector(&self) -> bool {
211        self.get(UPSTREAM_SOURCE_KEY)
212            .map(|s| {
213                s.eq_ignore_ascii_case(BATCH_POSIX_FS_CONNECTOR)
214                    || s.eq_ignore_ascii_case(ADBC_SNOWFLAKE_CONNECTOR)
215            })
216            .unwrap_or(false)
217    }
218
219    fn supports_full_reload_refresh(&self) -> bool {
220        self.get(UPSTREAM_SOURCE_KEY)
221            .map(|s| {
222                s.eq_ignore_ascii_case(OPENDAL_S3_CONNECTOR)
223                    || s.eq_ignore_ascii_case(GCS_CONNECTOR)
224                    || s.eq_ignore_ascii_case(BATCH_POSIX_FS_CONNECTOR)
225                    || s.eq_ignore_ascii_case(ICEBERG_CONNECTOR)
226                    || s.eq_ignore_ascii_case(ADBC_SNOWFLAKE_CONNECTOR)
227            })
228            .unwrap_or(false)
229    }
230
231    fn requires_singleton(&self) -> bool {
232        self.is_new_fs_connector() || self.is_iceberg_connector() || self.is_batch_connector()
233    }
234}
235
236impl<T: Get + GetKeyIter> WithPropertiesExt for T {}
237
238/// Options or properties extracted from the `WITH` clause of DDLs.
239#[derive(Default, Clone, Debug, PartialEq, Eq, Hash)]
240pub struct WithOptionsSecResolved {
241    inner: BTreeMap<String, String>,
242    secret_ref: BTreeMap<String, PbSecretRef>,
243}
244
245impl std::ops::Deref for WithOptionsSecResolved {
246    type Target = BTreeMap<String, String>;
247
248    fn deref(&self) -> &Self::Target {
249        &self.inner
250    }
251}
252
253impl std::ops::DerefMut for WithOptionsSecResolved {
254    fn deref_mut(&mut self) -> &mut Self::Target {
255        &mut self.inner
256    }
257}
258
259impl WithOptionsSecResolved {
260    /// Create a new [`WithOptions`] from a option [`BTreeMap`] and resolved secret ref.
261    pub fn new(inner: BTreeMap<String, String>, secret_ref: BTreeMap<String, PbSecretRef>) -> Self {
262        Self { inner, secret_ref }
263    }
264
265    pub fn as_plaintext(&self) -> &BTreeMap<String, String> {
266        &self.inner
267    }
268
269    pub fn as_secret(&self) -> &BTreeMap<String, PbSecretRef> {
270        &self.secret_ref
271    }
272
273    pub fn handle_update(
274        &mut self,
275        update_alter_props: BTreeMap<String, String>,
276        update_alter_secret_refs: BTreeMap<String, PbSecretRef>,
277    ) -> ConnectorResult<(Vec<SecretId>, Vec<SecretId>)> {
278        let old_secret_ids = self
279            .secret_ref
280            .values()
281            .map(|secret_ref| secret_ref.secret_id)
282            .collect::<BTreeSet<_>>();
283
284        // make sure the key in update_alter_props and update_alter_secret_refs not collide
285        for key in update_alter_props.keys() {
286            if update_alter_secret_refs.contains_key(key) {
287                return Err(
288                    anyhow::anyhow!("the key {} is set both in plaintext and secret", key).into(),
289                );
290            }
291        }
292
293        // remove legacy key if it's set in both plaintext and secret
294        // When a property changes from secret to plaintext, remove the old secret dependency
295        for k in update_alter_props.keys() {
296            self.secret_ref.remove(k);
297        }
298
299        // Handle secret ref updates
300        for k in update_alter_secret_refs.keys() {
301            // Remove any plaintext value for this key
302            self.inner.remove(k);
303        }
304
305        self.inner.extend(update_alter_props);
306        self.secret_ref.extend(update_alter_secret_refs);
307
308        let new_secret_ids = self
309            .secret_ref
310            .values()
311            .map(|secret_ref| secret_ref.secret_id)
312            .collect::<BTreeSet<_>>();
313        let to_add_secret_dep = new_secret_ids
314            .difference(&old_secret_ids)
315            .copied()
316            .collect();
317        let to_remove_secret_dep = old_secret_ids
318            .difference(&new_secret_ids)
319            .copied()
320            .collect();
321
322        Ok((to_add_secret_dep, to_remove_secret_dep))
323    }
324
325    /// Create a new [`WithOptions`] from a [`BTreeMap`].
326    pub fn without_secrets(inner: BTreeMap<String, String>) -> Self {
327        Self {
328            inner,
329            secret_ref: Default::default(),
330        }
331    }
332
333    /// Take the value of the option map and secret refs.
334    pub fn into_parts(self) -> (BTreeMap<String, String>, BTreeMap<String, PbSecretRef>) {
335        (self.inner, self.secret_ref)
336    }
337
338    pub fn value_eq_ignore_case(&self, key: &str, val: &str) -> bool {
339        if let Some(inner_val) = self.inner.get(key)
340            && inner_val.eq_ignore_ascii_case(val)
341        {
342            return true;
343        }
344        false
345    }
346}
347
348/// For `planner_test` crate so that it does not depend directly on `connector` crate just for `SinkFormatDesc`.
349impl TryFrom<&WithOptionsSecResolved> for Option<SinkFormatDesc> {
350    type Error = crate::sink::SinkError;
351
352    fn try_from(value: &WithOptionsSecResolved) -> std::result::Result<Self, Self::Error> {
353        let connector = value.get(crate::sink::CONNECTOR_TYPE_KEY);
354        let r#type = value.get(crate::sink::SINK_TYPE_OPTION);
355        match (connector, r#type) {
356            (Some(c), Some(t)) => SinkFormatDesc::from_legacy_type(c, t),
357            _ => Ok(None),
358        }
359    }
360}
361
362impl Get for WithOptionsSecResolved {
363    fn get(&self, key: &str) -> Option<&String> {
364        self.inner.get(key)
365    }
366}
367
368impl GetKeyIter for WithOptionsSecResolved {
369    fn key_iter(&self) -> impl Iterator<Item = &str> {
370        self.inner.keys().map(|s| s.as_str())
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    fn source_options(connector: &str) -> WithOptionsSecResolved {
379        WithOptionsSecResolved::without_secrets(BTreeMap::from([(
380            UPSTREAM_SOURCE_KEY.to_owned(),
381            connector.to_owned(),
382        )]))
383    }
384
385    #[test]
386    fn test_full_reload_refresh_connector_whitelist() {
387        for connector in [
388            OPENDAL_S3_CONNECTOR,
389            GCS_CONNECTOR,
390            BATCH_POSIX_FS_CONNECTOR,
391            ICEBERG_CONNECTOR,
392            ADBC_SNOWFLAKE_CONNECTOR,
393        ] {
394            assert!(
395                source_options(connector).supports_full_reload_refresh(),
396                "{connector} should support FULL_RELOAD refresh"
397            );
398        }
399
400        for connector in [KAFKA_CONNECTOR, POSIX_FS_CONNECTOR, AZBLOB_CONNECTOR] {
401            assert!(
402                !source_options(connector).supports_full_reload_refresh(),
403                "{connector} should not support FULL_RELOAD refresh"
404            );
405        }
406    }
407}