risingwave_connector/sink/encoder/
json.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;
16use std::sync::Arc;
17
18use anyhow::Context;
19use base64::Engine as _;
20use base64::engine::general_purpose;
21use chrono::{DateTime, Datelike, Timelike};
22use chrono_tz::Tz;
23use indexmap::IndexMap;
24use itertools::Itertools;
25use risingwave_common::array::{ArrayError, ArrayResult};
26use risingwave_common::catalog::{Field, Schema};
27use risingwave_common::row::Row;
28use risingwave_common::types::{DataType, DatumRef, JsonbVal, ScalarRefImpl, ToText};
29use risingwave_common::util::iter_util::ZipEqDebug;
30use serde_json::{Map, Value, json};
31use thiserror_ext::AsReport;
32
33use super::{
34    CustomJsonType, DateHandlingMode, JsonbHandlingMode, KafkaConnectParams, KafkaConnectParamsRef,
35    Result, RowEncoder, SerTo, TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
36};
37use crate::sink::SinkError;
38
39pub struct JsonEncoderConfig {
40    time_handling_mode: TimeHandlingMode,
41    date_handling_mode: DateHandlingMode,
42    timestamp_handling_mode: TimestampHandlingMode,
43    timestamptz_handling_mode: TimestamptzHandlingMode,
44    custom_json_type: CustomJsonType,
45    jsonb_handling_mode: JsonbHandlingMode,
46}
47
48pub struct JsonEncoder {
49    schema: Schema,
50    col_indices: Option<Vec<usize>>,
51    kafka_connect: Option<KafkaConnectParamsRef>,
52    config: JsonEncoderConfig,
53}
54
55impl JsonEncoder {
56    pub fn new(
57        schema: Schema,
58        col_indices: Option<Vec<usize>>,
59        date_handling_mode: DateHandlingMode,
60        timestamp_handling_mode: TimestampHandlingMode,
61        timestamptz_handling_mode: TimestamptzHandlingMode,
62        time_handling_mode: TimeHandlingMode,
63        jsonb_handling_mode: JsonbHandlingMode,
64    ) -> Self {
65        let config = JsonEncoderConfig {
66            time_handling_mode,
67            date_handling_mode,
68            timestamp_handling_mode,
69            timestamptz_handling_mode,
70            custom_json_type: CustomJsonType::None,
71            jsonb_handling_mode,
72        };
73        Self {
74            schema,
75            col_indices,
76            kafka_connect: None,
77            config,
78        }
79    }
80
81    pub fn new_with_es(schema: Schema, col_indices: Option<Vec<usize>>) -> Self {
82        let config = JsonEncoderConfig {
83            time_handling_mode: TimeHandlingMode::String,
84            date_handling_mode: DateHandlingMode::String,
85            timestamp_handling_mode: TimestampHandlingMode::String,
86            timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
87            custom_json_type: CustomJsonType::Es,
88            jsonb_handling_mode: JsonbHandlingMode::Dynamic,
89        };
90        Self {
91            schema,
92            col_indices,
93            kafka_connect: None,
94            config,
95        }
96    }
97
98    pub fn new_with_doris(
99        schema: Schema,
100        col_indices: Option<Vec<usize>>,
101        map: HashMap<String, u8>,
102    ) -> Self {
103        let config = JsonEncoderConfig {
104            time_handling_mode: TimeHandlingMode::Milli,
105            date_handling_mode: DateHandlingMode::String,
106            timestamp_handling_mode: TimestampHandlingMode::String,
107            timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
108            custom_json_type: CustomJsonType::Doris(map),
109            jsonb_handling_mode: JsonbHandlingMode::String,
110        };
111        Self {
112            schema,
113            col_indices,
114            kafka_connect: None,
115            config,
116        }
117    }
118
119    pub fn new_with_starrocks(
120        schema: Schema,
121        col_indices: Option<Vec<usize>>,
122        time_zone: Tz,
123    ) -> Self {
124        let config = JsonEncoderConfig {
125            time_handling_mode: TimeHandlingMode::Milli,
126            date_handling_mode: DateHandlingMode::String,
127            timestamp_handling_mode: TimestampHandlingMode::String,
128            timestamptz_handling_mode: TimestamptzHandlingMode::SpecifiedTimezoneWithoutSuffix(
129                time_zone,
130            ),
131            custom_json_type: CustomJsonType::StarRocks,
132            jsonb_handling_mode: JsonbHandlingMode::Dynamic,
133        };
134        Self {
135            schema,
136            col_indices,
137            kafka_connect: None,
138            config,
139        }
140    }
141
142    pub fn new_with_turbopuffer(schema: Schema, col_indices: Option<Vec<usize>>) -> Self {
143        let config = JsonEncoderConfig {
144            time_handling_mode: TimeHandlingMode::String,
145            date_handling_mode: DateHandlingMode::String,
146            timestamp_handling_mode: TimestampHandlingMode::Iso8601String,
147            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
148            custom_json_type: CustomJsonType::Turbopuffer,
149            jsonb_handling_mode: JsonbHandlingMode::Dynamic,
150        };
151        Self {
152            schema,
153            col_indices,
154            kafka_connect: None,
155            config,
156        }
157    }
158
159    pub fn with_kafka_connect(self, kafka_connect: KafkaConnectParams) -> Self {
160        Self {
161            kafka_connect: Some(Arc::new(kafka_connect)),
162            ..self
163        }
164    }
165}
166
167impl RowEncoder for JsonEncoder {
168    type Output = Map<String, Value>;
169
170    fn schema(&self) -> &Schema {
171        &self.schema
172    }
173
174    fn col_indices(&self) -> Option<&[usize]> {
175        self.col_indices.as_ref().map(Vec::as_ref)
176    }
177
178    fn encode_cols(
179        &self,
180        row: impl Row,
181        col_indices: impl Iterator<Item = usize>,
182    ) -> Result<Self::Output> {
183        let mut mappings = Map::with_capacity(self.schema.len());
184        let col_indices = col_indices.collect_vec();
185        for idx in &col_indices {
186            let field = &self.schema[*idx];
187            let key = field.name.clone();
188            let value = datum_to_json_object(field, row.datum_at(*idx), &self.config)
189                .map_err(|e| SinkError::Encode(e.to_report_string()))?;
190            mappings.insert(key, value);
191        }
192
193        Ok(if let Some(param) = &self.kafka_connect {
194            json_converter_with_schema(
195                Value::Object(mappings),
196                param.schema_name.clone(),
197                col_indices.into_iter().map(|i| &self.schema[i]),
198            )
199        } else {
200            mappings
201        })
202    }
203}
204
205impl SerTo<String> for Map<String, Value> {
206    fn ser_to(self) -> Result<String> {
207        Value::Object(self).ser_to()
208    }
209}
210
211impl SerTo<String> for Value {
212    fn ser_to(self) -> Result<String> {
213        Ok(self.to_string())
214    }
215}
216
217fn datum_to_json_object(
218    field: &Field,
219    datum: DatumRef<'_>,
220    config: &JsonEncoderConfig,
221) -> ArrayResult<Value> {
222    let scalar_ref = match datum {
223        None => {
224            return Ok(Value::Null);
225        }
226        Some(datum) => datum,
227    };
228
229    let data_type = field.data_type();
230
231    tracing::trace!("datum_to_json_object: {:?}, {:?}", data_type, scalar_ref);
232
233    let value = match (data_type, scalar_ref) {
234        (DataType::Boolean, ScalarRefImpl::Bool(v)) => {
235            json!(v)
236        }
237        (DataType::Int16, ScalarRefImpl::Int16(v)) => {
238            json!(v)
239        }
240        (DataType::Int32, ScalarRefImpl::Int32(v)) => {
241            json!(v)
242        }
243        (DataType::Int64, ScalarRefImpl::Int64(v)) => {
244            json!(v)
245        }
246        (DataType::Serial, ScalarRefImpl::Serial(v)) => {
247            if matches!(&config.custom_json_type, CustomJsonType::Turbopuffer) {
248                json!(v.into_inner())
249            } else {
250                // The serial type needs to be handled as a string to prevent primary key conflicts caused by the precision issues of JSON numbers.
251                json!(format!("{:#018x}", v.into_inner()))
252            }
253        }
254        (DataType::Float32, ScalarRefImpl::Float32(v)) => {
255            json!(f32::from(v))
256        }
257        (DataType::Float64, ScalarRefImpl::Float64(v)) => {
258            json!(f64::from(v))
259        }
260        (DataType::Varchar, ScalarRefImpl::Utf8(v)) => {
261            json!(v)
262        }
263        // Doris/Starrocks will convert out-of-bounds decimal and -INF, INF, NAN to NULL
264        (DataType::Decimal, ScalarRefImpl::Decimal(mut v)) => match &config.custom_json_type {
265            CustomJsonType::Doris(map) => {
266                let s = map.get(&field.name).unwrap();
267                v.rescale(*s as u32);
268                json!(v.to_text())
269            }
270            CustomJsonType::Turbopuffer => {
271                let value = f64::try_from(v).map_err(|err| {
272                    ArrayError::internal(format!(
273                        "failed to convert decimal to f64: {}",
274                        err.as_report()
275                    ))
276                })?;
277                serde_json::Number::from_f64(value)
278                    .map(Value::Number)
279                    .ok_or_else(|| {
280                        ArrayError::internal(
281                            "failed to encode non-finite decimal as JSON number".to_owned(),
282                        )
283                    })?
284            }
285            CustomJsonType::Es | CustomJsonType::None | CustomJsonType::StarRocks => {
286                json!(v.to_text())
287            }
288        },
289        (DataType::Timestamptz, ScalarRefImpl::Timestamptz(v)) => {
290            match config.timestamptz_handling_mode {
291                TimestamptzHandlingMode::UtcString => {
292                    let parsed = v.to_datetime_utc();
293                    let v = parsed.to_rfc3339_opts(chrono::SecondsFormat::Micros, true);
294                    json!(v)
295                }
296                TimestamptzHandlingMode::UtcWithoutSuffix => {
297                    let parsed = v.to_datetime_utc().naive_utc();
298                    let v = parsed.format("%Y-%m-%d %H:%M:%S%.6f").to_string();
299                    json!(v)
300                }
301                TimestamptzHandlingMode::SpecifiedTimezoneWithoutSuffix(time_zone) => {
302                    let parsed = v.to_datetime_in_zone(time_zone).naive_local();
303                    let v = parsed.format("%Y-%m-%d %H:%M:%S%.6f").to_string();
304                    json!(v)
305                }
306                TimestamptzHandlingMode::Micro => json!(v.timestamp_micros()),
307                TimestamptzHandlingMode::Milli => json!(v.timestamp_millis()),
308            }
309        }
310        (DataType::Time, ScalarRefImpl::Time(v)) => match config.time_handling_mode {
311            TimeHandlingMode::Milli => {
312                // todo: just ignore the nanos part to avoid leap second complex
313                json!(v.0.num_seconds_from_midnight() as i64 * 1000)
314            }
315            TimeHandlingMode::String => {
316                let a = v.0.format("%H:%M:%S%.6f").to_string();
317                json!(a)
318            }
319        },
320        (DataType::Date, ScalarRefImpl::Date(v)) => match config.date_handling_mode {
321            DateHandlingMode::FromCe => json!(v.0.num_days_from_ce()),
322            DateHandlingMode::FromEpoch => {
323                let duration = v.0 - DateTime::UNIX_EPOCH.date_naive();
324                json!(duration.num_days())
325            }
326            DateHandlingMode::String => {
327                let a = v.0.format("%Y-%m-%d").to_string();
328                json!(a)
329            }
330        },
331        (DataType::Timestamp, ScalarRefImpl::Timestamp(v)) => {
332            match config.timestamp_handling_mode {
333                TimestampHandlingMode::Milli => json!(v.0.and_utc().timestamp_millis()),
334                TimestampHandlingMode::String => {
335                    json!(v.0.format("%Y-%m-%d %H:%M:%S%.6f").to_string())
336                }
337                TimestampHandlingMode::Iso8601String => {
338                    json!(v.0.format("%Y-%m-%dT%H:%M:%S%.6f").to_string())
339                }
340            }
341        }
342        (DataType::Bytea, ScalarRefImpl::Bytea(v)) => {
343            json!(general_purpose::STANDARD.encode(v))
344        }
345        // P<years>Y<months>M<days>DT<hours>H<minutes>M<seconds>S
346        (DataType::Interval, ScalarRefImpl::Interval(v)) => {
347            json!(v.as_iso_8601())
348        }
349
350        (DataType::Jsonb, ScalarRefImpl::Jsonb(jsonb_ref)) => match config.jsonb_handling_mode {
351            JsonbHandlingMode::String => {
352                json!(jsonb_ref.to_string())
353            }
354            JsonbHandlingMode::Dynamic => JsonbVal::from(jsonb_ref).take(),
355        },
356        (DataType::List(lt), ScalarRefImpl::List(list_ref)) => {
357            let elems = list_ref.iter();
358            let mut vec = Vec::with_capacity(elems.len());
359            let inner_field = Field::unnamed(lt.into_elem());
360            for sub_datum_ref in elems {
361                let value = datum_to_json_object(&inner_field, sub_datum_ref, config)?;
362                vec.push(value);
363            }
364            json!(vec)
365        }
366        (DataType::Vector(_), ScalarRefImpl::Vector(vector)) => {
367            let elems = vector.as_raw_slice();
368            let mut vec = Vec::with_capacity(elems.len());
369            for v in elems {
370                let value = serde_json::Number::from_f64(*v as _)
371                    .map(Value::Number)
372                    .unwrap_or(Value::Null);
373                vec.push(value);
374            }
375            json!(vec)
376        }
377        (DataType::Struct(st), ScalarRefImpl::Struct(struct_ref)) => {
378            match config.custom_json_type {
379                CustomJsonType::Doris(_) => {
380                    // We need to ensure that the order of elements in the json matches the insertion order.
381                    let mut map = IndexMap::with_capacity(st.len());
382                    for (sub_datum_ref, sub_field) in struct_ref.iter_fields_ref().zip_eq_debug(
383                        st.iter()
384                            .map(|(name, dt)| Field::with_name(dt.clone(), name)),
385                    ) {
386                        let value = datum_to_json_object(&sub_field, sub_datum_ref, config)?;
387                        map.insert(sub_field.name.clone(), value);
388                    }
389                    Value::String(
390                        serde_json::to_string(&map).context("failed to serialize into JSON")?,
391                    )
392                }
393                CustomJsonType::StarRocks => {
394                    return Err(ArrayError::internal(
395                        "starrocks can't support struct".to_owned(),
396                    ));
397                }
398                CustomJsonType::Es | CustomJsonType::None | CustomJsonType::Turbopuffer => {
399                    let mut map = Map::with_capacity(st.len());
400                    for (sub_datum_ref, sub_field) in struct_ref.iter_fields_ref().zip_eq_debug(
401                        st.iter()
402                            .map(|(name, dt)| Field::with_name(dt.clone(), name)),
403                    ) {
404                        let value = datum_to_json_object(&sub_field, sub_datum_ref, config)?;
405                        map.insert(sub_field.name.clone(), value);
406                    }
407                    json!(map)
408                }
409            }
410        }
411        // TODO(map): support map
412        (data_type, scalar_ref) => {
413            return Err(ArrayError::internal(format!(
414                "datum_to_json_object: unsupported data type: field name: {:?}, logical type: {:?}, physical type: {:?}",
415                field.name, data_type, scalar_ref
416            )));
417        }
418    };
419
420    Ok(value)
421}
422
423fn json_converter_with_schema<'a>(
424    object: Value,
425    name: String,
426    fields: impl Iterator<Item = &'a Field>,
427) -> Map<String, Value> {
428    let mut mapping = Map::with_capacity(2);
429    mapping.insert(
430        "schema".to_owned(),
431        json!({
432            "type": "struct",
433            "fields": fields.map(|field| {
434                let mut mapping = type_as_json_schema(&field.data_type);
435                mapping.insert("field".to_owned(), json!(field.name));
436                mapping
437            }).collect_vec(),
438            "optional": false,
439            "name": name,
440        }),
441    );
442    mapping.insert("payload".to_owned(), object);
443    mapping
444}
445
446// reference: https://github.com/apache/kafka/blob/80982c4ae3fe6be127b48ec09caff11ab5f87c69/connect/json/src/main/java/org/apache/kafka/connect/json/JsonSchema.java#L39
447pub(crate) fn schema_type_mapping(rw_type: &DataType) -> &'static str {
448    match rw_type {
449        DataType::Boolean => "boolean",
450        DataType::Int16 => "int16",
451        DataType::Int32 => "int32",
452        DataType::Int64 => "int64",
453        DataType::Float32 => "float",
454        DataType::Float64 => "double",
455        DataType::Decimal => "string",
456        DataType::Date => "int32",
457        DataType::Varchar => "string",
458        DataType::Time => "int64",
459        DataType::Timestamp => "int64",
460        DataType::Timestamptz => "string",
461        DataType::Interval => "string",
462        DataType::Struct(_) => "struct",
463        DataType::List(_) => "array",
464        DataType::Vector(_) => "array",
465        DataType::Bytea => "bytes",
466        DataType::Jsonb => "string",
467        DataType::Serial => "string",
468        DataType::Int256 => "string",
469        DataType::Map(_) => "map",
470    }
471}
472
473fn type_as_json_schema(rw_type: &DataType) -> Map<String, Value> {
474    let mut mapping = Map::with_capacity(4); // type + optional + fields/items + field
475    mapping.insert("type".to_owned(), json!(schema_type_mapping(rw_type)));
476    mapping.insert("optional".to_owned(), json!(true));
477    match rw_type {
478        DataType::Struct(struct_type) => {
479            let sub_fields = struct_type
480                .iter()
481                .map(|(sub_name, sub_type)| {
482                    let mut sub_mapping = type_as_json_schema(sub_type);
483                    sub_mapping.insert("field".to_owned(), json!(sub_name));
484                    sub_mapping
485                })
486                .collect_vec();
487            mapping.insert("fields".to_owned(), json!(sub_fields));
488        }
489        DataType::List(list_type) => {
490            mapping.insert(
491                "items".to_owned(),
492                json!(type_as_json_schema(list_type.elem())),
493            );
494        }
495        _ => {}
496    }
497
498    mapping
499}
500
501#[cfg(test)]
502mod tests {
503    use risingwave_common::row::OwnedRow;
504    use risingwave_common::types::{
505        Date, Decimal, Interval, Scalar, ScalarImpl, StructRef, StructType, StructValue, Time,
506        Timestamp,
507    };
508
509    use super::*;
510
511    #[test]
512    fn test_starrocks_timestamptz_encoding() {
513        let schema = Schema::new(vec![Field::with_name(DataType::Timestamptz, "ts")]);
514        let encoder = JsonEncoder::new_with_starrocks(schema, None, chrono_tz::Asia::Shanghai);
515        let tstz = "2018-01-26T18:30:09.453Z".parse().unwrap();
516        let row = OwnedRow::new(vec![Some(ScalarImpl::Timestamptz(tstz))]);
517
518        let encoded = encoder.encode(row).unwrap();
519
520        assert_eq!(
521            encoded.get("ts"),
522            Some(&json!("2018-01-27 02:30:09.453000"))
523        );
524    }
525
526    #[test]
527    fn test_to_json_basic_type() {
528        let mock_field = Field {
529            data_type: DataType::Boolean,
530            name: Default::default(),
531        };
532
533        let config = JsonEncoderConfig {
534            time_handling_mode: TimeHandlingMode::Milli,
535            date_handling_mode: DateHandlingMode::FromCe,
536            timestamp_handling_mode: TimestampHandlingMode::String,
537            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
538            custom_json_type: CustomJsonType::None,
539            jsonb_handling_mode: JsonbHandlingMode::String,
540        };
541
542        let boolean_value = datum_to_json_object(
543            &Field {
544                data_type: DataType::Boolean,
545                ..mock_field.clone()
546            },
547            Some(ScalarImpl::Bool(false).as_scalar_ref_impl()),
548            &config,
549        )
550        .unwrap();
551        assert_eq!(boolean_value, json!(false));
552
553        let int16_value = datum_to_json_object(
554            &Field {
555                data_type: DataType::Int16,
556                ..mock_field.clone()
557            },
558            Some(ScalarImpl::Int16(16).as_scalar_ref_impl()),
559            &config,
560        )
561        .unwrap();
562        assert_eq!(int16_value, json!(16));
563
564        let int64_value = datum_to_json_object(
565            &Field {
566                data_type: DataType::Int64,
567                ..mock_field.clone()
568            },
569            Some(ScalarImpl::Int64(i64::MAX).as_scalar_ref_impl()),
570            &config,
571        )
572        .unwrap();
573        assert_eq!(
574            serde_json::to_string(&int64_value).unwrap(),
575            i64::MAX.to_string()
576        );
577
578        let serial_value = datum_to_json_object(
579            &Field {
580                data_type: DataType::Serial,
581                ..mock_field.clone()
582            },
583            Some(ScalarImpl::Serial(i64::MAX.into()).as_scalar_ref_impl()),
584            &config,
585        )
586        .unwrap();
587        assert_eq!(
588            serde_json::to_string(&serial_value).unwrap(),
589            format!("\"{:#018x}\"", i64::MAX)
590        );
591
592        let turbopuffer_config = JsonEncoderConfig {
593            time_handling_mode: TimeHandlingMode::String,
594            date_handling_mode: DateHandlingMode::String,
595            timestamp_handling_mode: TimestampHandlingMode::String,
596            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
597            custom_json_type: CustomJsonType::Turbopuffer,
598            jsonb_handling_mode: JsonbHandlingMode::Dynamic,
599        };
600        let turbopuffer_serial_value = datum_to_json_object(
601            &Field {
602                data_type: DataType::Serial,
603                ..mock_field.clone()
604            },
605            Some(ScalarImpl::Serial(i64::MAX.into()).as_scalar_ref_impl()),
606            &turbopuffer_config,
607        )
608        .unwrap();
609        assert_eq!(turbopuffer_serial_value, json!(i64::MAX));
610        let turbopuffer_decimal_value = datum_to_json_object(
611            &Field {
612                data_type: DataType::Decimal,
613                ..mock_field.clone()
614            },
615            Some(ScalarImpl::Decimal(Decimal::try_from(1.25).unwrap()).as_scalar_ref_impl()),
616            &turbopuffer_config,
617        )
618        .unwrap();
619        assert_eq!(turbopuffer_decimal_value, json!(1.25));
620        assert!(
621            datum_to_json_object(
622                &Field {
623                    data_type: DataType::Decimal,
624                    ..mock_field.clone()
625                },
626                Some(ScalarImpl::Decimal(Decimal::NaN).as_scalar_ref_impl()),
627                &turbopuffer_config,
628            )
629            .unwrap_err()
630            .to_string()
631            .contains("non-finite decimal")
632        );
633
634        // https://github.com/debezium/debezium/blob/main/debezium-core/src/main/java/io/debezium/time/ZonedTimestamp.java
635        let tstz_inner = "2018-01-26T18:30:09.453Z".parse().unwrap();
636        let tstz_value = datum_to_json_object(
637            &Field {
638                data_type: DataType::Timestamptz,
639                ..mock_field.clone()
640            },
641            Some(ScalarImpl::Timestamptz(tstz_inner).as_scalar_ref_impl()),
642            &config,
643        )
644        .unwrap();
645        assert_eq!(tstz_value, "2018-01-26T18:30:09.453000Z");
646
647        let unix_wo_suffix_config = JsonEncoderConfig {
648            time_handling_mode: TimeHandlingMode::Milli,
649            date_handling_mode: DateHandlingMode::FromCe,
650            timestamp_handling_mode: TimestampHandlingMode::String,
651            timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
652            custom_json_type: CustomJsonType::None,
653            jsonb_handling_mode: JsonbHandlingMode::String,
654        };
655
656        let tstz_inner = "2018-01-26T18:30:09.453Z".parse().unwrap();
657        let tstz_value = datum_to_json_object(
658            &Field {
659                data_type: DataType::Timestamptz,
660                ..mock_field.clone()
661            },
662            Some(ScalarImpl::Timestamptz(tstz_inner).as_scalar_ref_impl()),
663            &unix_wo_suffix_config,
664        )
665        .unwrap();
666        assert_eq!(tstz_value, "2018-01-26 18:30:09.453000");
667
668        let timestamp_milli_config = JsonEncoderConfig {
669            time_handling_mode: TimeHandlingMode::String,
670            date_handling_mode: DateHandlingMode::FromCe,
671            timestamp_handling_mode: TimestampHandlingMode::Milli,
672            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
673            custom_json_type: CustomJsonType::None,
674            jsonb_handling_mode: JsonbHandlingMode::String,
675        };
676        let ts_value = datum_to_json_object(
677            &Field {
678                data_type: DataType::Timestamp,
679                ..mock_field.clone()
680            },
681            Some(
682                ScalarImpl::Timestamp(Timestamp::from_timestamp_uncheck(1000, 0))
683                    .as_scalar_ref_impl(),
684            ),
685            &timestamp_milli_config,
686        )
687        .unwrap();
688        assert_eq!(ts_value, json!(1000 * 1000));
689
690        let ts_value = datum_to_json_object(
691            &Field {
692                data_type: DataType::Timestamp,
693                ..mock_field.clone()
694            },
695            Some(
696                ScalarImpl::Timestamp(Timestamp::from_timestamp_uncheck(1000, 0))
697                    .as_scalar_ref_impl(),
698            ),
699            &config,
700        )
701        .unwrap();
702        assert_eq!(ts_value, json!("1970-01-01 00:16:40.000000".to_owned()));
703
704        // Represents the number of milliseconds past midnigh, org.apache.kafka.connect.data.Time
705        let time_value = datum_to_json_object(
706            &Field {
707                data_type: DataType::Time,
708                ..mock_field.clone()
709            },
710            Some(
711                ScalarImpl::Time(Time::from_num_seconds_from_midnight_uncheck(1000, 0))
712                    .as_scalar_ref_impl(),
713            ),
714            &config,
715        )
716        .unwrap();
717        assert_eq!(time_value, json!(1000 * 1000));
718
719        let interval_value = datum_to_json_object(
720            &Field {
721                data_type: DataType::Interval,
722                ..mock_field.clone()
723            },
724            Some(
725                ScalarImpl::Interval(Interval::from_month_day_usec(13, 2, 1000000))
726                    .as_scalar_ref_impl(),
727            ),
728            &config,
729        )
730        .unwrap();
731        assert_eq!(interval_value, json!("P1Y1M2DT0H0M1S"));
732
733        let mut map = HashMap::default();
734        map.insert("aaa".to_owned(), 5_u8);
735        let doris_config = JsonEncoderConfig {
736            time_handling_mode: TimeHandlingMode::String,
737            date_handling_mode: DateHandlingMode::String,
738            timestamp_handling_mode: TimestampHandlingMode::String,
739            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
740            custom_json_type: CustomJsonType::Doris(map),
741            jsonb_handling_mode: JsonbHandlingMode::String,
742        };
743        let decimal = datum_to_json_object(
744            &Field {
745                data_type: DataType::Decimal,
746                name: "aaa".to_owned(),
747            },
748            Some(ScalarImpl::Decimal(Decimal::try_from(1.1111111).unwrap()).as_scalar_ref_impl()),
749            &doris_config,
750        )
751        .unwrap();
752        assert_eq!(decimal, json!("1.11111"));
753
754        let date_value = datum_to_json_object(
755            &Field {
756                data_type: DataType::Date,
757                ..mock_field.clone()
758            },
759            Some(ScalarImpl::Date(Date::from_ymd_uncheck(1970, 1, 1)).as_scalar_ref_impl()),
760            &config,
761        )
762        .unwrap();
763        assert_eq!(date_value, json!(719163));
764
765        let from_epoch_config = JsonEncoderConfig {
766            time_handling_mode: TimeHandlingMode::String,
767            date_handling_mode: DateHandlingMode::FromEpoch,
768            timestamp_handling_mode: TimestampHandlingMode::String,
769            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
770            custom_json_type: CustomJsonType::None,
771            jsonb_handling_mode: JsonbHandlingMode::String,
772        };
773        let date_value = datum_to_json_object(
774            &Field {
775                data_type: DataType::Date,
776                ..mock_field.clone()
777            },
778            Some(ScalarImpl::Date(Date::from_ymd_uncheck(1970, 1, 1)).as_scalar_ref_impl()),
779            &from_epoch_config,
780        )
781        .unwrap();
782        assert_eq!(date_value, json!(0));
783
784        let doris_config = JsonEncoderConfig {
785            time_handling_mode: TimeHandlingMode::String,
786            date_handling_mode: DateHandlingMode::String,
787            timestamp_handling_mode: TimestampHandlingMode::String,
788            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
789            custom_json_type: CustomJsonType::Doris(HashMap::default()),
790            jsonb_handling_mode: JsonbHandlingMode::String,
791        };
792        let date_value = datum_to_json_object(
793            &Field {
794                data_type: DataType::Date,
795                ..mock_field.clone()
796            },
797            Some(ScalarImpl::Date(Date::from_ymd_uncheck(2010, 10, 10)).as_scalar_ref_impl()),
798            &doris_config,
799        )
800        .unwrap();
801        assert_eq!(date_value, json!("2010-10-10"));
802
803        let value = StructValue::new(vec![
804            Some(3_i32.to_scalar_value()),
805            Some(2_i32.to_scalar_value()),
806            Some(1_i32.to_scalar_value()),
807        ]);
808
809        let interval_value = datum_to_json_object(
810            &Field {
811                data_type: DataType::Struct(StructType::new(vec![
812                    ("v3", DataType::Int32),
813                    ("v2", DataType::Int32),
814                    ("v1", DataType::Int32),
815                ])),
816                ..mock_field.clone()
817            },
818            Some(ScalarRefImpl::Struct(StructRef::ValueRef { val: &value })),
819            &doris_config,
820        )
821        .unwrap();
822        assert_eq!(interval_value, json!("{\"v3\":3,\"v2\":2,\"v1\":1}"));
823
824        let encode_jsonb_obj_config = JsonEncoderConfig {
825            time_handling_mode: TimeHandlingMode::String,
826            date_handling_mode: DateHandlingMode::String,
827            timestamp_handling_mode: TimestampHandlingMode::String,
828            timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
829            custom_json_type: CustomJsonType::None,
830            jsonb_handling_mode: JsonbHandlingMode::Dynamic,
831        };
832        let json_value = datum_to_json_object(
833            &Field {
834                data_type: DataType::Jsonb,
835                ..mock_field
836            },
837            Some(ScalarImpl::Jsonb(JsonbVal::from(json!([1, 2, 3]))).as_scalar_ref_impl()),
838            &encode_jsonb_obj_config,
839        )
840        .unwrap();
841        assert_eq!(json_value, json!([1, 2, 3]));
842    }
843
844    #[test]
845    fn test_generate_json_converter_schema() {
846        let fields = vec![
847            Field {
848                data_type: DataType::Boolean,
849                name: "v1".into(),
850            },
851            Field {
852                data_type: DataType::Int16,
853                name: "v2".into(),
854            },
855            Field {
856                data_type: DataType::Int32,
857                name: "v3".into(),
858            },
859            Field {
860                data_type: DataType::Float32,
861                name: "v4".into(),
862            },
863            Field {
864                data_type: DataType::Decimal,
865                name: "v5".into(),
866            },
867            Field {
868                data_type: DataType::Date,
869                name: "v6".into(),
870            },
871            Field {
872                data_type: DataType::Varchar,
873                name: "v7".into(),
874            },
875            Field {
876                data_type: DataType::Time,
877                name: "v8".into(),
878            },
879            Field {
880                data_type: DataType::Interval,
881                name: "v9".into(),
882            },
883            Field {
884                data_type: DataType::Struct(StructType::new(vec![
885                    ("a", DataType::Timestamp),
886                    ("b", DataType::Timestamptz),
887                    (
888                        "c",
889                        DataType::Struct(StructType::new(vec![
890                            ("aa", DataType::Int64),
891                            ("bb", DataType::Float64),
892                        ])),
893                    ),
894                ])),
895                name: "v10".into(),
896            },
897            Field {
898                data_type: DataType::list(DataType::list(DataType::Struct(StructType::new(vec![
899                    ("aa", DataType::Int64),
900                    ("bb", DataType::Float64),
901                ])))),
902                name: "v11".into(),
903            },
904            Field {
905                data_type: DataType::Jsonb,
906                name: "12".into(),
907            },
908            Field {
909                data_type: DataType::Serial,
910                name: "13".into(),
911            },
912            Field {
913                data_type: DataType::Int256,
914                name: "14".into(),
915            },
916        ];
917        let schema =
918            json_converter_with_schema(json!({}), "test".to_owned(), fields.iter())["schema"]
919                .to_string();
920        let ans = r#"{"fields":[{"field":"v1","optional":true,"type":"boolean"},{"field":"v2","optional":true,"type":"int16"},{"field":"v3","optional":true,"type":"int32"},{"field":"v4","optional":true,"type":"float"},{"field":"v5","optional":true,"type":"string"},{"field":"v6","optional":true,"type":"int32"},{"field":"v7","optional":true,"type":"string"},{"field":"v8","optional":true,"type":"int64"},{"field":"v9","optional":true,"type":"string"},{"field":"v10","fields":[{"field":"a","optional":true,"type":"int64"},{"field":"b","optional":true,"type":"string"},{"field":"c","fields":[{"field":"aa","optional":true,"type":"int64"},{"field":"bb","optional":true,"type":"double"}],"optional":true,"type":"struct"}],"optional":true,"type":"struct"},{"field":"v11","items":{"items":{"fields":[{"field":"aa","optional":true,"type":"int64"},{"field":"bb","optional":true,"type":"double"}],"optional":true,"type":"struct"},"optional":true,"type":"array"},"optional":true,"type":"array"},{"field":"12","optional":true,"type":"string"},{"field":"13","optional":true,"type":"string"},{"field":"14","optional":true,"type":"string"}],"name":"test","optional":false,"type":"struct"}"#;
921        assert_eq!(schema, ans);
922    }
923}
risingwave_connector/sink/encoder/json.rs

risingwave_connector/sink/encoder/
json.rs
