risingwave_connector/sink/encoder/
mod.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, HashMap, HashSet};
16use std::sync::Arc;
17
18use chrono_tz::Tz;
19use risingwave_common::catalog::Schema;
20use risingwave_common::row::Row;
21
22use crate::sink::Result;
23
24mod avro;
25mod bson;
26pub mod bytes;
27mod json;
28mod proto;
29pub mod template;
30pub mod text;
31
32pub use avro::{AvroEncoder, AvroHeader};
33pub use bson::BsonEncoder;
34pub use json::JsonEncoder;
35pub use proto::{ProtoEncoder, ProtoHeader};
36
37/// Encode a row of a relation into
38/// * an object in json
39/// * a message in protobuf
40/// * a record in avro
41pub trait RowEncoder {
42    type Output: SerTo<Vec<u8>>;
43
44    fn encode_cols(
45        &self,
46        row: impl Row,
47        col_indices: impl Iterator<Item = usize>,
48    ) -> Result<Self::Output>;
49    fn schema(&self) -> &Schema;
50    fn col_indices(&self) -> Option<&[usize]>;
51
52    fn encode(&self, row: impl Row) -> Result<Self::Output> {
53        assert_eq!(row.len(), self.schema().len());
54        match self.col_indices() {
55            Some(col_indices) => self.encode_cols(row, col_indices.iter().copied()),
56            None => self.encode_cols(row, 0..self.schema().len()),
57        }
58    }
59}
60
61/// Do the actual encoding from
62/// * an json object
63/// * a protobuf message
64/// * an avro record
65///   into
66/// * string (required by kinesis key)
67/// * bytes
68///
69/// This is like `TryInto` but allows us to `impl<T: SerTo<String>> SerTo<Vec<u8>> for T`.
70///
71/// Note that `serde` does not fit here because its data model does not contain logical types.
72/// For example, although `chrono::DateTime` implements `Serialize`,
73/// it produces avro String rather than avro `TimestampMicros`.
74pub trait SerTo<T> {
75    fn ser_to(self) -> Result<T>;
76}
77
78impl<T: SerTo<String>> SerTo<Vec<u8>> for T {
79    fn ser_to(self) -> Result<Vec<u8>> {
80        self.ser_to().map(|s: String| s.into_bytes())
81    }
82}
83
84impl<T> SerTo<T> for T {
85    fn ser_to(self) -> Result<T> {
86        Ok(self)
87    }
88}
89
90#[derive(Clone, Copy, Default)]
91pub enum DateHandlingMode {
92    #[default]
93    FromCe,
94    FromEpoch,
95    String,
96}
97
98impl DateHandlingMode {
99    pub const OPTION_KEY: &'static str = "date.handling.mode";
100
101    pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
102        match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
103            // Preserve the legacy encoding when the option is absent.
104            Some("from_ce") | None => Ok(Self::FromCe),
105            Some("from_epoch") => Ok(Self::FromEpoch),
106            Some("string") => Ok(Self::String),
107            Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
108                "unrecognized {} value {}",
109                Self::OPTION_KEY,
110                v
111            ))),
112        }
113    }
114}
115
116/// Useful for both json and protobuf
117#[derive(Clone, Copy)]
118pub enum TimestampHandlingMode {
119    Milli,
120    String,
121    /// ISO 8601 string without a timezone suffix.
122    Iso8601String,
123}
124
125#[derive(Clone, Copy)]
126pub enum TimeHandlingMode {
127    Milli,
128    String,
129}
130
131#[derive(Clone, Copy, Default)]
132pub enum TimestamptzHandlingMode {
133    #[default]
134    UtcString,
135    UtcWithoutSuffix,
136    SpecifiedTimezoneWithoutSuffix(Tz),
137    Micro,
138    Milli,
139}
140
141impl TimestamptzHandlingMode {
142    pub const FRONTEND_DEFAULT: &'static str = "utc_string";
143    pub const OPTION_KEY: &'static str = "timestamptz.handling.mode";
144
145    pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
146        match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
147            Some(Self::FRONTEND_DEFAULT) => Ok(Self::UtcString),
148            Some("utc_without_suffix") => Ok(Self::UtcWithoutSuffix),
149            Some("micro") => Ok(Self::Micro),
150            Some("milli") => Ok(Self::Milli),
151            Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
152                "unrecognized {} value {}",
153                Self::OPTION_KEY,
154                v
155            ))),
156            // This is not a good default. We just have to select it when no option is provided
157            // for compatibility with old version.
158            None => Ok(Self::UtcWithoutSuffix),
159        }
160    }
161}
162
163#[derive(Clone)]
164pub struct DorisJsonConfig {
165    pub decimal_scale: HashMap<String, u8>,
166    pub variant_columns: HashSet<String>,
167}
168
169#[derive(Clone)]
170pub enum CustomJsonType {
171    // Doris's json need date is string.
172    // The internal order of the struct should follow the insertion order.
173    // The decimal needs verification and calibration.
174    Doris(DorisJsonConfig),
175    // Es's json need jsonb is struct
176    Es,
177    // starrocks' need jsonb is struct
178    StarRocks,
179    // turbopuffer expects serial and decimal attributes to match `int` and `float` schema types.
180    Turbopuffer,
181    None,
182}
183
184/// How the jsonb type is encoded.
185///
186/// - `String`: encode jsonb as string. `[1, true, "foo"] -> "[1, true, \"foo\"]"`
187/// - `Dynamic`: encode jsonb as json type dynamically. `[1, true, "foo"] -> [1, true, "foo"]`
188#[derive(Clone, Copy)]
189pub enum JsonbHandlingMode {
190    String,
191    Dynamic,
192}
193
194impl JsonbHandlingMode {
195    pub const OPTION_KEY: &'static str = "jsonb.handling.mode";
196
197    pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
198        match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
199            Some("string") | None => Ok(Self::String),
200            Some("dynamic") => Ok(Self::Dynamic),
201            Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
202                "unrecognized {} value {}",
203                Self::OPTION_KEY,
204                v
205            ))),
206        }
207    }
208}
209
210#[derive(Debug)]
211struct FieldEncodeError {
212    message: String,
213    rev_path: Vec<String>,
214}
215
216impl FieldEncodeError {
217    fn new(message: impl std::fmt::Display) -> Self {
218        Self {
219            message: message.to_string(),
220            rev_path: vec![],
221        }
222    }
223
224    fn with_name(mut self, name: &str) -> Self {
225        self.rev_path.push(name.into());
226        self
227    }
228}
229
230impl std::fmt::Display for FieldEncodeError {
231    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
232        use itertools::Itertools;
233
234        write!(
235            f,
236            "encode '{}' error: {}",
237            self.rev_path.iter().rev().join("."),
238            self.message
239        )
240    }
241}
242
243impl From<FieldEncodeError> for super::SinkError {
244    fn from(value: FieldEncodeError) -> Self {
245        Self::Encode(value.to_string())
246    }
247}
248
249#[derive(Clone)]
250pub struct KafkaConnectParams {
251    pub schema_name: String,
252}
253
254type KafkaConnectParamsRef = Arc<KafkaConnectParams>;
risingwave_connector/sink/encoder/mod.rs

risingwave_connector/sink/encoder/
mod.rs
