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
98/// Useful for both json and protobuf
99#[derive(Clone, Copy)]
100pub enum TimestampHandlingMode {
101    Milli,
102    String,
103    /// ISO 8601 string without a timezone suffix.
104    Iso8601String,
105}
106
107#[derive(Clone, Copy)]
108pub enum TimeHandlingMode {
109    Milli,
110    String,
111}
112
113#[derive(Clone, Copy, Default)]
114pub enum TimestamptzHandlingMode {
115    #[default]
116    UtcString,
117    UtcWithoutSuffix,
118    SpecifiedTimezoneWithoutSuffix(Tz),
119    Micro,
120    Milli,
121}
122
123impl TimestamptzHandlingMode {
124    pub const FRONTEND_DEFAULT: &'static str = "utc_string";
125    pub const OPTION_KEY: &'static str = "timestamptz.handling.mode";
126
127    pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
128        match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
129            Some(Self::FRONTEND_DEFAULT) => Ok(Self::UtcString),
130            Some("utc_without_suffix") => Ok(Self::UtcWithoutSuffix),
131            Some("micro") => Ok(Self::Micro),
132            Some("milli") => Ok(Self::Milli),
133            Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
134                "unrecognized {} value {}",
135                Self::OPTION_KEY,
136                v
137            ))),
138            // This is not a good default. We just have to select it when no option is provided
139            // for compatibility with old version.
140            None => Ok(Self::UtcWithoutSuffix),
141        }
142    }
143}
144
145#[derive(Clone)]
146pub struct DorisJsonConfig {
147    pub decimal_scale: HashMap<String, u8>,
148    pub variant_columns: HashSet<String>,
149}
150
151#[derive(Clone)]
152pub enum CustomJsonType {
153    // Doris's json need date is string.
154    // The internal order of the struct should follow the insertion order.
155    // The decimal needs verification and calibration.
156    Doris(DorisJsonConfig),
157    // Es's json need jsonb is struct
158    Es,
159    // starrocks' need jsonb is struct
160    StarRocks,
161    // turbopuffer expects serial and decimal attributes to match `int` and `float` schema types.
162    Turbopuffer,
163    None,
164}
165
166/// How the jsonb type is encoded.
167///
168/// - `String`: encode jsonb as string. `[1, true, "foo"] -> "[1, true, \"foo\"]"`
169/// - `Dynamic`: encode jsonb as json type dynamically. `[1, true, "foo"] -> [1, true, "foo"]`
170#[derive(Clone, Copy)]
171pub enum JsonbHandlingMode {
172    String,
173    Dynamic,
174}
175
176impl JsonbHandlingMode {
177    pub const OPTION_KEY: &'static str = "jsonb.handling.mode";
178
179    pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
180        match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
181            Some("string") | None => Ok(Self::String),
182            Some("dynamic") => Ok(Self::Dynamic),
183            Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
184                "unrecognized {} value {}",
185                Self::OPTION_KEY,
186                v
187            ))),
188        }
189    }
190}
191
192#[derive(Debug)]
193struct FieldEncodeError {
194    message: String,
195    rev_path: Vec<String>,
196}
197
198impl FieldEncodeError {
199    fn new(message: impl std::fmt::Display) -> Self {
200        Self {
201            message: message.to_string(),
202            rev_path: vec![],
203        }
204    }
205
206    fn with_name(mut self, name: &str) -> Self {
207        self.rev_path.push(name.into());
208        self
209    }
210}
211
212impl std::fmt::Display for FieldEncodeError {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        use itertools::Itertools;
215
216        write!(
217            f,
218            "encode '{}' error: {}",
219            self.rev_path.iter().rev().join("."),
220            self.message
221        )
222    }
223}
224
225impl From<FieldEncodeError> for super::SinkError {
226    fn from(value: FieldEncodeError) -> Self {
227        Self::Encode(value.to_string())
228    }
229}
230
231#[derive(Clone)]
232pub struct KafkaConnectParams {
233    pub schema_name: String,
234}
235
236type KafkaConnectParamsRef = Arc<KafkaConnectParams>;
risingwave_connector/sink/encoder/mod.rs

risingwave_connector/sink/encoder/
mod.rs
