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

risingwave_connector/sink/encoder/
json.rs
