risingwave_connector/parser/
utils.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::collections::BTreeMap;
16
17use anyhow::Context;
18use bytes::Bytes;
19use reqwest::Url;
20use risingwave_common::bail;
21use risingwave_common::types::{Datum, DatumCow, DatumRef, ScalarRefImpl};
22use risingwave_pb::data::DataType as PbDataType;
23
24use crate::aws_utils::load_file_descriptor_from_s3;
25use crate::connector_common::AwsAuthProps;
26use crate::error::ConnectorResult;
27use crate::source::SourceMeta;
28
29macro_rules! log_error {
30    ($name:expr, $err:expr, $message:expr) => {
31        if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
32            tracing::error!(
33                column = $name,
34                error = %$err.as_report(),
35                suppressed_count,
36                $message,
37            );
38        }
39    };
40}
41pub(crate) use log_error;
42use risingwave_pb::plan_common::additional_column;
43use risingwave_pb::plan_common::additional_column::ColumnType;
44
45use crate::parser::{AccessError, AccessResult};
46use crate::source::cdc::DebeziumCdcMeta;
47
48fn get_topic<'a>(
49    props: &'a BTreeMap<String, String>,
50    connector_topic_key: &str,
51) -> ConnectorResult<&'a String> {
52    const TOPIC_KEY: &str = "topic";
53
54    if let Some(topic) = props.get(connector_topic_key) {
55        return Ok(topic);
56    }
57    if let Some(topic) = props.get(TOPIC_KEY) {
58        return Ok(topic);
59    }
60
61    bail!("Must specify '{}' or '{}'", connector_topic_key, TOPIC_KEY)
62}
63
64/// Get the Kafka topic name.
65pub(super) fn get_kafka_topic(props: &BTreeMap<String, String>) -> ConnectorResult<&String> {
66    get_topic(props, "kafka.topic")
67}
68
69/// Get the Pulsar topic name.
70pub(super) fn get_pulsar_topic(props: &BTreeMap<String, String>) -> ConnectorResult<&String> {
71    get_topic(props, "pulsar.topic")
72}
73
74/// download bytes from http(s) url
75pub(super) async fn download_from_http(location: &Url) -> ConnectorResult<Bytes> {
76    let res = reqwest::get(location.clone())
77        .await
78        .with_context(|| format!("failed to make request to {location}"))?
79        .error_for_status()
80        .with_context(|| format!("http request failed for {location}"))?;
81
82    let bytes = res
83        .bytes()
84        .await
85        .with_context(|| format!("failed to read HTTP body of {location}"))?;
86
87    Ok(bytes)
88}
89
90// For parser that doesn't support key currently
91#[macro_export]
92macro_rules! only_parse_payload {
93    ($self:ident, $payload:ident, $writer:ident) => {
94        if let Some(payload) = $payload {
95            $self.parse_inner(payload, $writer).await
96        } else {
97            risingwave_common::bail!("empty payload with non-empty key")
98        }
99    };
100}
101
102/// Load raw bytes from:
103/// * local file, for on-premise or testing.
104/// * http/https, for common usage.
105/// * s3 file location format: <s3://bucket_name/file_name>
106pub(super) async fn bytes_from_url(
107    url: &Url,
108    config: Option<&AwsAuthProps>,
109) -> ConnectorResult<Vec<u8>> {
110    match (url.scheme(), config) {
111        // TODO(Tao): support local file only when it's compiled in debug mode.
112        ("file", _) => {
113            let path = url
114                .to_file_path()
115                .ok()
116                .with_context(|| format!("illegal path: {url}"))?;
117            Ok(std::fs::read(&path)
118                .with_context(|| format!("failed to read file from `{}`", path.display()))?)
119        }
120        ("https" | "http", _) => Ok(download_from_http(url).await?.into()),
121        ("s3", Some(config)) => load_file_descriptor_from_s3(url, config).await,
122        (scheme, _) => bail!("path scheme `{scheme}` is not supported"),
123    }
124}
125
126pub fn extract_timestamp_from_meta(meta: &SourceMeta) -> DatumRef<'_> {
127    match meta {
128        SourceMeta::Kafka(kafka_meta) => kafka_meta.extract_timestamp(),
129        SourceMeta::DebeziumCdc(cdc_meta) => cdc_meta.extract_timestamp(),
130        SourceMeta::Kinesis(kinesis_meta) => kinesis_meta.extract_timestamp(),
131        _ => None,
132    }
133}
134
135pub fn extract_cdc_meta_column<'a>(
136    cdc_meta: &'a DebeziumCdcMeta,
137    column_type: &additional_column::ColumnType,
138    column_name: &str,
139) -> AccessResult<DatumRef<'a>> {
140    match column_type {
141        ColumnType::Timestamp(_) => Ok(cdc_meta.extract_timestamp()),
142        ColumnType::DatabaseName(_) => Ok(cdc_meta.extract_database_name()),
143        // `table_name` in additional columns should be the object name only, not the routing key
144        // (which can be `schema.table` or `db.table` depending on the connector).
145        ColumnType::TableName(_) => Ok(cdc_meta.extract_table_name_only()),
146        _ => Err(AccessError::UnsupportedAdditionalColumn {
147            name: column_name.to_owned(),
148        }),
149    }
150}
151
152pub fn extract_headers_from_meta(meta: &SourceMeta) -> Option<Datum> {
153    match meta {
154        SourceMeta::Kafka(kafka_meta) => kafka_meta.extract_headers(), /* expect output of type `array[struct<varchar, bytea>]` */
155        SourceMeta::Pulsar(pulsar_meta) => pulsar_meta.extract_headers(), /* expect output of type `array[struct<varchar, bytea>]` */
156        SourceMeta::Nats(nats_meta) => nats_meta.extract_headers(), /* expect output of type `array[struct<varchar, bytea>]` */
157        _ => None,
158    }
159}
160
161pub fn extract_pulsar_message_id_data_from_meta(meta: &SourceMeta) -> Option<DatumRef<'_>> {
162    match meta {
163        SourceMeta::Pulsar(pulsar_meta) => pulsar_meta
164            .ack_message_id
165            .as_ref()
166            .map(|message_id| Some(ScalarRefImpl::Bytea(message_id))),
167        _ => None,
168    }
169}
170
171pub fn extract_header_inner_from_meta<'a>(
172    meta: &'a SourceMeta,
173    inner_field: &str,
174    data_type: Option<&PbDataType>,
175) -> Option<DatumCow<'a>> {
176    match meta {
177        SourceMeta::Kafka(kafka_meta) => kafka_meta.extract_header_inner(inner_field, data_type), /* expect output of type `bytea` or `varchar` */
178        SourceMeta::Pulsar(pulsar_meta) => pulsar_meta.extract_header_inner(inner_field, data_type),
179        SourceMeta::Nats(nats_meta) => nats_meta.extract_header_inner(inner_field, data_type), /* expect output of type `bytea` or `varchar` */
180        _ => None,
181    }
182}
183
184pub fn extract_subject_from_meta(meta: &SourceMeta) -> Option<DatumRef<'_>> {
185    match meta {
186        SourceMeta::Nats(nats_meta) => Some(nats_meta.extract_subject()),
187        _ => None,
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use async_nats::HeaderMap;
194    use risingwave_common::types::{DataType, ScalarImpl, ToOwnedDatum};
195
196    use super::*;
197    use crate::source::nats::source::NatsMeta;
198
199    fn nats_meta(subject: &str, headers: Option<HeaderMap>) -> SourceMeta {
200        SourceMeta::Nats(NatsMeta {
201            subject: subject.to_owned(),
202            headers,
203        })
204    }
205
206    #[test]
207    fn extract_headers_from_meta_dispatches_to_nats() {
208        let mut headers = HeaderMap::new();
209        headers.insert("k", "v");
210        let meta = nats_meta("subj", Some(headers));
211
212        let datum = extract_headers_from_meta(&meta).expect("outer Option");
213        assert!(matches!(datum, Some(ScalarImpl::List(_))));
214    }
215
216    #[test]
217    fn extract_header_inner_from_meta_dispatches_to_nats() {
218        let mut headers = HeaderMap::new();
219        headers.insert("trace", "abc");
220        let meta = nats_meta("subj", Some(headers));
221        let pb = DataType::Varchar.to_protobuf();
222
223        let result = extract_header_inner_from_meta(&meta, "trace", Some(&pb))
224            .unwrap()
225            .to_owned_datum();
226        assert_eq!(result, Some(ScalarImpl::Utf8("abc".into())));
227    }
228}
risingwave_connector/parser/utils.rs

risingwave_connector/parser/
utils.rs
