risingwave_connector/parser/
additional_columns.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::{HashMap, HashSet};
16use std::sync::LazyLock;
17
18use risingwave_common::bail;
19use risingwave_common::catalog::{ColumnCatalog, ColumnDesc, ColumnId, max_column_id};
20use risingwave_common::types::{DataType, StructType};
21use risingwave_pb::plan_common::additional_column::ColumnType as AdditionalColumnType;
22use risingwave_pb::plan_common::{
23    AdditionalCollectionName, AdditionalColumn, AdditionalColumnFilename, AdditionalColumnHeader,
24    AdditionalColumnHeaders, AdditionalColumnKey, AdditionalColumnOffset,
25    AdditionalColumnPartition, AdditionalColumnPayload, AdditionalColumnPulsarMessageIdData,
26    AdditionalColumnTimestamp, AdditionalDatabaseName, AdditionalSchemaName, AdditionalSubject,
27    AdditionalTableName,
28};
29
30use crate::dispatch_cdc_source_type;
31use crate::error::ConnectorResult;
32use crate::source::cdc::{CdcSourceDdlValidator, CdcSourceType};
33use crate::source::{
34    AZBLOB_CONNECTOR, GCS_CONNECTOR, KAFKA_CONNECTOR, KINESIS_CONNECTOR, MQTT_CONNECTOR,
35    NATS_CONNECTOR, OPENDAL_S3_CONNECTOR, POSIX_FS_CONNECTOR, PULSAR_CONNECTOR,
36};
37
38// Hidden additional columns connectors which do not support `include` syntax.
39pub static COMMON_COMPATIBLE_ADDITIONAL_COLUMNS: LazyLock<HashSet<&'static str>> =
40    LazyLock::new(|| HashSet::from(["partition", "offset"]));
41
42pub static COMPATIBLE_ADDITIONAL_COLUMNS: LazyLock<HashMap<&'static str, HashSet<&'static str>>> =
43    LazyLock::new(|| {
44        HashMap::from([
45            (
46                KAFKA_CONNECTOR,
47                HashSet::from([
48                    "key",
49                    "timestamp",
50                    "partition",
51                    "offset",
52                    "header",
53                    "payload",
54                ]),
55            ),
56            (
57                PULSAR_CONNECTOR,
58                HashSet::from([
59                    "key",
60                    "partition",
61                    "offset",
62                    "payload",
63                    "message_id_data",
64                    "header",
65                ]),
66            ),
67            (
68                KINESIS_CONNECTOR,
69                HashSet::from(["key", "partition", "offset", "timestamp", "payload"]),
70            ),
71            (
72                NATS_CONNECTOR,
73                HashSet::from(["partition", "offset", "payload", "subject", "header"]),
74            ),
75            (
76                OPENDAL_S3_CONNECTOR,
77                HashSet::from(["file", "offset", "payload"]),
78            ),
79            (GCS_CONNECTOR, HashSet::from(["file", "offset", "payload"])),
80            (
81                AZBLOB_CONNECTOR,
82                HashSet::from(["file", "offset", "payload"]),
83            ),
84            (
85                POSIX_FS_CONNECTOR,
86                HashSet::from(["file", "offset", "payload"]),
87            ),
88            (MQTT_CONNECTOR, HashSet::from(["offset", "partition"])),
89        ])
90    });
91
92// For CDC backfill table, the additional columns are added to the schema of `StreamCdcScan`
93pub static CDC_BACKFILL_TABLE_ADDITIONAL_COLUMNS: LazyLock<Option<HashSet<&'static str>>> =
94    LazyLock::new(|| {
95        Some(HashSet::from([
96            "timestamp",
97            "database_name",
98            "schema_name",
99            "table_name",
100        ]))
101    });
102
103pub fn gen_default_addition_col_name(
104    connector_name: &str,
105    additional_col_type: &str,
106    inner_field_name: Option<&str>,
107    data_type: Option<&DataType>,
108) -> String {
109    let legacy_dt_name = data_type.map(|dt| format!("{:?}", dt).to_lowercase());
110    let col_name = [
111        Some(connector_name),
112        Some(additional_col_type),
113        inner_field_name,
114        legacy_dt_name.as_deref(),
115    ];
116    col_name.iter().fold("_rw".to_owned(), |name, ele| {
117        if let Some(ele) = ele {
118            format!("{}_{}", name, ele)
119        } else {
120            name
121        }
122    })
123}
124
125pub fn build_additional_column_desc(
126    column_id: ColumnId,
127    connector_name: &str,
128    additional_col_type: &str,
129    column_alias: Option<String>,
130    inner_field_name: Option<&str>,
131    data_type: Option<&DataType>,
132    reject_unknown_connector: bool,
133    is_cdc_backfill_table: bool,
134) -> ConnectorResult<ColumnDesc> {
135    let checked_by_cdc_source = reject_unknown_connector
136        && !is_cdc_backfill_table
137        && dispatch_cdc_source_type!(
138            CdcSourceType::from(connector_name),
139            T,
140            {
141                T::check_include_column(additional_col_type)?;
142                true
143            },
144            false
145        );
146    if !checked_by_cdc_source {
147        let supported_columns = if is_cdc_backfill_table {
148            CDC_BACKFILL_TABLE_ADDITIONAL_COLUMNS.as_ref()
149        } else {
150            COMPATIBLE_ADDITIONAL_COLUMNS.get(connector_name)
151        };
152        let compatible_columns = match (supported_columns, reject_unknown_connector) {
153            (Some(compat_cols), _) => compat_cols,
154            (None, false) => &COMMON_COMPATIBLE_ADDITIONAL_COLUMNS,
155            (None, true) => {
156                bail!(
157                    "additional column is not supported for connector {}, acceptable connectors: {:?}",
158                    connector_name,
159                    COMPATIBLE_ADDITIONAL_COLUMNS.keys(),
160                );
161            }
162        };
163        if !compatible_columns.contains(additional_col_type) {
164            bail!(
165                "additional column type {} is not supported for connector {}, acceptable column types: {:?}",
166                additional_col_type,
167                connector_name,
168                compatible_columns
169            );
170        }
171    }
172
173    let column_name = column_alias.unwrap_or_else(|| {
174        gen_default_addition_col_name(
175            connector_name,
176            additional_col_type,
177            inner_field_name,
178            data_type,
179        )
180    });
181
182    let col_desc = match additional_col_type {
183        "key" => ColumnDesc::named_with_additional_column(
184            column_name,
185            column_id,
186            DataType::Bytea,
187            AdditionalColumn {
188                column_type: Some(AdditionalColumnType::Key(AdditionalColumnKey {})),
189            },
190        ),
191
192        "timestamp" => ColumnDesc::named_with_additional_column(
193            column_name,
194            column_id,
195            DataType::Timestamptz,
196            AdditionalColumn {
197                column_type: Some(AdditionalColumnType::Timestamp(
198                    AdditionalColumnTimestamp {},
199                )),
200            },
201        ),
202        "partition" => ColumnDesc::named_with_additional_column(
203            column_name,
204            column_id,
205            DataType::Varchar,
206            AdditionalColumn {
207                column_type: Some(AdditionalColumnType::Partition(
208                    AdditionalColumnPartition {},
209                )),
210            },
211        ),
212        "payload" => ColumnDesc::named_with_additional_column(
213            column_name,
214            column_id,
215            DataType::Jsonb,
216            AdditionalColumn {
217                column_type: Some(AdditionalColumnType::Payload(AdditionalColumnPayload {})),
218            },
219        ),
220        "offset" => ColumnDesc::named_with_additional_column(
221            column_name,
222            column_id,
223            DataType::Varchar,
224            AdditionalColumn {
225                column_type: Some(AdditionalColumnType::Offset(AdditionalColumnOffset {})),
226            },
227        ),
228
229        "file" => ColumnDesc::named_with_additional_column(
230            column_name,
231            column_id,
232            DataType::Varchar,
233            AdditionalColumn {
234                column_type: Some(AdditionalColumnType::Filename(AdditionalColumnFilename {})),
235            },
236        ),
237        "header" => build_header_catalog(column_id, &column_name, inner_field_name, data_type),
238        "database_name" => ColumnDesc::named_with_additional_column(
239            column_name,
240            column_id,
241            DataType::Varchar,
242            AdditionalColumn {
243                column_type: Some(AdditionalColumnType::DatabaseName(
244                    AdditionalDatabaseName {},
245                )),
246            },
247        ),
248        "schema_name" => ColumnDesc::named_with_additional_column(
249            column_name,
250            column_id,
251            DataType::Varchar,
252            AdditionalColumn {
253                column_type: Some(AdditionalColumnType::SchemaName(AdditionalSchemaName {})),
254            },
255        ),
256        "table_name" => ColumnDesc::named_with_additional_column(
257            column_name,
258            column_id,
259            DataType::Varchar,
260            AdditionalColumn {
261                column_type: Some(AdditionalColumnType::TableName(AdditionalTableName {})),
262            },
263        ),
264        "collection_name" => ColumnDesc::named_with_additional_column(
265            column_name,
266            column_id,
267            DataType::Varchar,
268            AdditionalColumn {
269                column_type: Some(AdditionalColumnType::CollectionName(
270                    AdditionalCollectionName {},
271                )),
272            },
273        ),
274        "subject" => ColumnDesc::named_with_additional_column(
275            column_name,
276            column_id,
277            DataType::Varchar, // Assuming subject is a string
278            AdditionalColumn {
279                column_type: Some(AdditionalColumnType::Subject(AdditionalSubject {})),
280            },
281        ),
282        "message_id_data" => ColumnDesc::named_with_additional_column(
283            column_name,
284            column_id,
285            DataType::Bytea,
286            AdditionalColumn {
287                column_type: Some(AdditionalColumnType::PulsarMessageIdData(
288                    AdditionalColumnPulsarMessageIdData {},
289                )),
290            },
291        ),
292        _ => unreachable!(),
293    };
294
295    Ok(col_desc)
296}
297
298pub fn derive_pulsar_message_id_data_column(
299    connector_name: &str,
300    column_exist: &mut Vec<bool>,
301    additional_columns: &mut Vec<ColumnDesc>,
302) {
303    // additional columns already check the max_column_id
304    // so we can take the max column id of additional columns as the max column id of all columns
305    let max_column_id = additional_columns
306        .iter()
307        .fold(ColumnId::first_user_column(), |a, b| a.max(b.column_id));
308
309    // assume user does not include `message_id_data` column
310    column_exist.push(false);
311    additional_columns.push(
312        build_additional_column_desc(
313            max_column_id.next(),
314            connector_name,
315            "message_id_data",
316            None,
317            None,
318            None,
319            false,
320            false,
321        )
322        .unwrap(),
323    );
324}
325
326/// Utility function for adding partition and offset columns to the columns, if not specified by the user.
327///
328/// ## Returns
329/// - `columns_exist`: whether 1. `partition`/`file` and 2. `offset` columns are included in `columns`.
330/// - `additional_columns`: The `ColumnCatalog` for `partition`/`file` and `offset` columns.
331pub fn source_add_partition_offset_cols(
332    columns: &[ColumnCatalog],
333    connector_name: &str,
334    skip_col_id: bool,
335) -> (Vec<bool>, Vec<ColumnDesc>) {
336    let mut columns_exist = vec![false; 2];
337
338    let mut last_column_id = max_column_id(columns);
339    let mut assign_col_id = || {
340        if skip_col_id {
341            // col id will be filled outside later. Here just use a placeholder.
342            ColumnId::placeholder()
343        } else {
344            last_column_id = last_column_id.next();
345            last_column_id
346        }
347    };
348
349    let additional_columns: Vec<_> = {
350        let compat_col_types = COMPATIBLE_ADDITIONAL_COLUMNS
351            .get(connector_name)
352            .unwrap_or(&COMMON_COMPATIBLE_ADDITIONAL_COLUMNS);
353        ["partition", "file", "offset"]
354            .iter()
355            .filter_map(|col_type| {
356                if compat_col_types.contains(col_type) {
357                    Some(
358                        build_additional_column_desc(
359                            assign_col_id(),
360                            connector_name,
361                            col_type,
362                            None,
363                            None,
364                            None,
365                            false,
366                            false,
367                        )
368                        .unwrap(),
369                    )
370                } else {
371                    None
372                }
373            })
374            .collect()
375    };
376    assert_eq!(additional_columns.len(), 2);
377    use risingwave_pb::plan_common::additional_column::ColumnType;
378    assert_matches::assert_matches!(
379        additional_columns[0].additional_column,
380        AdditionalColumn {
381            column_type: Some(ColumnType::Partition(_) | ColumnType::Filename(_)),
382        }
383    );
384    assert_matches::assert_matches!(
385        additional_columns[1].additional_column,
386        AdditionalColumn {
387            column_type: Some(ColumnType::Offset(_)),
388        }
389    );
390
391    // Check if partition/file/offset columns are included explicitly.
392    for col in columns {
393        match col.column_desc.additional_column {
394            AdditionalColumn {
395                column_type: Some(ColumnType::Partition(_) | ColumnType::Filename(_)),
396            } => {
397                columns_exist[0] = true;
398            }
399            AdditionalColumn {
400                column_type: Some(ColumnType::Offset(_)),
401            } => {
402                columns_exist[1] = true;
403            }
404            _ => (),
405        }
406    }
407
408    (columns_exist, additional_columns)
409}
410
411fn build_header_catalog(
412    column_id: ColumnId,
413    col_name: &str,
414    inner_field_name: Option<&str>,
415    data_type: Option<&DataType>,
416) -> ColumnDesc {
417    if let Some(inner) = inner_field_name {
418        let data_type = data_type.unwrap_or(&DataType::Bytea);
419        let pb_data_type = data_type.to_protobuf();
420        ColumnDesc::named_with_additional_column(
421            col_name,
422            column_id,
423            data_type.clone(),
424            AdditionalColumn {
425                column_type: Some(AdditionalColumnType::HeaderInner(AdditionalColumnHeader {
426                    inner_field: inner.to_owned(),
427                    data_type: Some(pb_data_type),
428                })),
429            },
430        )
431    } else {
432        ColumnDesc::named_with_additional_column(
433            col_name,
434            column_id,
435            DataType::list(get_kafka_header_item_datatype()),
436            AdditionalColumn {
437                column_type: Some(AdditionalColumnType::Headers(AdditionalColumnHeaders {})),
438            },
439        )
440    }
441}
442
443pub fn get_kafka_header_item_datatype() -> DataType {
444    let struct_inner = vec![("key", DataType::Varchar), ("value", DataType::Bytea)];
445    DataType::Struct(StructType::new(struct_inner))
446}
447
448#[cfg(test)]
449mod test {
450    use risingwave_pb::plan_common::additional_column::ColumnType as AdditionalColumnType;
451
452    use super::*;
453
454    #[test]
455    fn test_gen_default_addition_col_name() {
456        assert_eq!(
457            gen_default_addition_col_name("kafka", "key", None, None),
458            "_rw_kafka_key"
459        );
460        assert_eq!(
461            gen_default_addition_col_name("kafka", "header", Some("inner"), None),
462            "_rw_kafka_header_inner"
463        );
464        assert_eq!(
465            gen_default_addition_col_name(
466                "kafka",
467                "header",
468                Some("inner"),
469                Some(&DataType::Varchar)
470            ),
471            "_rw_kafka_header_inner_varchar"
472        );
473    }
474
475    #[test]
476    fn test_build_pulsar_header_additional_column() {
477        let col = build_additional_column_desc(
478            ColumnId::new(1),
479            PULSAR_CONNECTOR,
480            "header",
481            Some("pulsar_header".to_owned()),
482            Some("tenant"),
483            Some(&DataType::Varchar),
484            true,
485            false,
486        )
487        .unwrap();
488
489        assert_eq!(col.name, "pulsar_header");
490        assert_eq!(col.data_type, DataType::Varchar);
491        assert_matches::assert_matches!(
492            col.additional_column.column_type,
493            Some(AdditionalColumnType::HeaderInner(ref header))
494                if header.inner_field == "tenant"
495        );
496    }
497}
risingwave_connector/parser/additional_columns.rs

risingwave_connector/parser/
additional_columns.rs
