feat(sink): support encode jsonb data as dynamic json type in sink (#…

…17693)

src/connector/src/sink/encoder/json.rs

@@ -30,8 +30,8 @@ use serde_json::{json, Map, Value};
 use thiserror_ext::AsReport;
 
 use super::{
- CustomJsonType, DateHandlingMode, KafkaConnectParams, KafkaConnectParamsRef, Result,
- RowEncoder, SerTo, TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
+ CustomJsonType, DateHandlingMode, JsonbHandlingMode, KafkaConnectParams, KafkaConnectParamsRef,
+ Result, RowEncoder, SerTo, TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
 };
 use crate::sink::SinkError;
 
@@ -41,6 +41,7 @@ pub struct JsonEncoderConfig {
  timestamp_handling_mode: TimestampHandlingMode,
  timestamptz_handling_mode: TimestamptzHandlingMode,
  custom_json_type: CustomJsonType,
+ jsonb_handling_mode: JsonbHandlingMode,
 }
 
 pub struct JsonEncoder {
@@ -58,13 +59,15 @@ impl JsonEncoder {
  timestamp_handling_mode: TimestampHandlingMode,
  timestamptz_handling_mode: TimestamptzHandlingMode,
  time_handling_mode: TimeHandlingMode,
+ jsonb_handling_mode: JsonbHandlingMode,
  ) -> Self {
  let config = JsonEncoderConfig {
  time_handling_mode,
  date_handling_mode,
  timestamp_handling_mode,
  timestamptz_handling_mode,
  custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode,
  };
  Self {
  schema,
@@ -81,6 +84,7 @@ impl JsonEncoder {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
  custom_json_type: CustomJsonType::Es,
+ jsonb_handling_mode: JsonbHandlingMode::Dynamic,
  };
  Self {
  schema,
@@ -101,6 +105,7 @@ impl JsonEncoder {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
  custom_json_type: CustomJsonType::Doris(map),
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
  Self {
  schema,
@@ -117,6 +122,7 @@ impl JsonEncoder {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
  custom_json_type: CustomJsonType::StarRocks,
+ jsonb_handling_mode: JsonbHandlingMode::Dynamic,
  };
  Self {
  schema,
@@ -289,11 +295,12 @@ fn datum_to_json_object(
  (DataType::Interval, ScalarRefImpl::Interval(v)) => {
  json!(v.as_iso_8601())
  }
- (DataType::Jsonb, ScalarRefImpl::Jsonb(jsonb_ref)) => match &config.custom_json_type {
-  CustomJsonType::Es | CustomJsonType::StarRocks => JsonbVal::from(jsonb_ref).take(),
- CustomJsonType::Doris(_) | CustomJsonType::None => {
+
+ (DataType::Jsonb, ScalarRefImpl::Jsonb(jsonb_ref)) => match config.jsonb_handling_mode {
+ JsonbHandlingMode::String => {
  json!(jsonb_ref.to_string())
  }
+ JsonbHandlingMode::Dynamic => JsonbVal::from(jsonb_ref).take(),
  },
  (DataType::List(datatype), ScalarRefImpl::List(list_ref)) => {
  let elems = list_ref.iter();
@@ -446,6 +453,7 @@ mod tests {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
  custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
 
  let boolean_value = datum_to_json_object(
@@ -517,6 +525,7 @@ mod tests {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcWithoutSuffix,
  custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
 
  let tstz_inner = "2018-01-26T18:30:09.453Z".parse().unwrap();
@@ -537,6 +546,7 @@ mod tests {
  timestamp_handling_mode: TimestampHandlingMode::Milli,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
  custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
  let ts_value = datum_to_json_object(
  &Field {
@@ -603,6 +613,7 @@ mod tests {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
  custom_json_type: CustomJsonType::Doris(map),
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
  let decimal = datum_to_json_object(
  &Field {
@@ -628,11 +639,12 @@ mod tests {
  assert_eq!(date_value, json!(719163));
 
  let from_epoch_config = JsonEncoderConfig {
- time_handling_mode: TimeHandlingMode::Milli,
+ time_handling_mode: TimeHandlingMode::String,
  date_handling_mode: DateHandlingMode::FromEpoch,
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
  custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
  let date_value = datum_to_json_object(
  &Field {
@@ -651,6 +663,7 @@ mod tests {
  timestamp_handling_mode: TimestampHandlingMode::String,
  timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
  custom_json_type: CustomJsonType::Doris(HashMap::default()),
+ jsonb_handling_mode: JsonbHandlingMode::String,
  };
  let date_value = datum_to_json_object(
  &Field {
@@ -683,6 +696,25 @@ mod tests {
  )
  .unwrap();
  assert_eq!(interval_value, json!("{\"v3\":3,\"v2\":2,\"v1\":1}"));
+
+ let encode_jsonb_obj_config = JsonEncoderConfig {
+ time_handling_mode: TimeHandlingMode::String,
+ date_handling_mode: DateHandlingMode::String,
+ timestamp_handling_mode: TimestampHandlingMode::String,
+ timestamptz_handling_mode: TimestamptzHandlingMode::UtcString,
+ custom_json_type: CustomJsonType::None,
+ jsonb_handling_mode: JsonbHandlingMode::Dynamic,
+ };
+ let json_value = datum_to_json_object(
+ &Field {
+ data_type: DataType::Jsonb,
+ ..mock_field.clone()
+ },
+ Some(ScalarImpl::Jsonb(JsonbVal::from(json!([1, 2, 3]))).as_scalar_ref_impl()),
+ &encode_jsonb_obj_config,
+ )
+ .unwrap();
+ assert_eq!(json_value, json!([1, 2, 3]));
  }
 
  #[test]

src/connector/src/sink/encoder/mod.rs

@@ -148,6 +148,31 @@ pub enum CustomJsonType {
  None,
 }
 
+/// How the jsonb type is encoded.
+///
+/// - `String`: encode jsonb as string. `[1, true, "foo"] -> "[1, true, \"foo\"]"`
+/// - `Dynamic`: encode jsonb as json type dynamically. `[1, true, "foo"] -> [1, true, "foo"]`
+pub enum JsonbHandlingMode {
+ String,
+ Dynamic,
+}
+
+impl JsonbHandlingMode {
+ pub const OPTION_KEY: &'static str = "jsonb.handling.mode";
+
+ pub fn from_options(options: &BTreeMap<String, String>) -> Result<Self> {
+ match options.get(Self::OPTION_KEY).map(std::ops::Deref::deref) {
+ Some("string") | None => Ok(Self::String),
+ Some("dynamic") => Ok(Self::Dynamic),
+ Some(v) => Err(super::SinkError::Config(anyhow::anyhow!(
+ "unrecognized {} value {}",
+ Self::OPTION_KEY,
+ v
+ ))),
+ }
+ }
+}
+
 #[derive(Debug)]
 struct FieldEncodeError {
  message: String,

src/connector/src/sink/formatter/debezium_json.rs

@@ -21,8 +21,8 @@ use tracing::warn;
 
 use super::{Result, SinkFormatter, StreamChunk};
 use crate::sink::encoder::{
- DateHandlingMode, JsonEncoder, RowEncoder, TimeHandlingMode, TimestampHandlingMode,
- TimestamptzHandlingMode,
+ DateHandlingMode, JsonEncoder, JsonbHandlingMode, RowEncoder, TimeHandlingMode,
+ TimestampHandlingMode, TimestamptzHandlingMode,
 };
 use crate::tri;
 
@@ -69,6 +69,7 @@ impl DebeziumJsonFormatter {
  TimestampHandlingMode::Milli,
  TimestamptzHandlingMode::UtcString,
  TimeHandlingMode::Milli,
+ JsonbHandlingMode::String,
  );
  let val_encoder = JsonEncoder::new(
  schema.clone(),
@@ -77,6 +78,7 @@ impl DebeziumJsonFormatter {
  TimestampHandlingMode::Milli,
  TimestamptzHandlingMode::UtcString,
  TimeHandlingMode::Milli,
+ JsonbHandlingMode::String,
  );
  Self {
  schema,
@@ -397,6 +399,7 @@ mod tests {
  TimestampHandlingMode::Milli,
  TimestamptzHandlingMode::UtcString,
  TimeHandlingMode::Milli,
+ JsonbHandlingMode::String,
  );
  let json_chunk = chunk_to_json(chunk, &encoder).unwrap();
  let schema_json = schema_to_json(&schema, "test_db", "test_table");

src/connector/src/sink/formatter/mod.rs

@@ -31,7 +31,8 @@ use super::catalog::{SinkEncode, SinkFormat, SinkFormatDesc};
 use super::encoder::template::TemplateEncoder;
 use super::encoder::text::TextEncoder;
 use super::encoder::{
- DateHandlingMode, KafkaConnectParams, TimeHandlingMode, TimestamptzHandlingMode,
+ DateHandlingMode, JsonbHandlingMode, KafkaConnectParams, TimeHandlingMode,
+ TimestamptzHandlingMode,
 };
 use super::redis::{KEY_FORMAT, VALUE_FORMAT};
 use crate::sink::encoder::{
@@ -113,13 +114,15 @@ pub trait EncoderBuild: Sized {
 impl EncoderBuild for JsonEncoder {
  async fn build(b: EncoderParams<'_>, pk_indices: Option<Vec<usize>>) -> Result<Self> {
  let timestamptz_mode = TimestamptzHandlingMode::from_options(&b.format_desc.options)?;
+ let jsonb_handling_mode = JsonbHandlingMode::from_options(&b.format_desc.options)?;
  let encoder = JsonEncoder::new(
  b.schema,
  pk_indices,
  DateHandlingMode::FromCe,
  TimestampHandlingMode::Milli,
  timestamptz_mode,
  TimeHandlingMode::Milli,
+ jsonb_handling_mode,
  );
  let encoder = if let Some(s) = b.format_desc.options.get("schemas.enable") {
  match s.to_lowercase().parse::<bool>() {

src/connector/src/sink/kafka.rs

@@ -590,7 +590,7 @@ mod test {
 
  use super::*;
  use crate::sink::encoder::{
- DateHandlingMode, JsonEncoder, TimeHandlingMode, TimestampHandlingMode,
+ DateHandlingMode, JsonEncoder, JsonbHandlingMode, TimeHandlingMode, TimestampHandlingMode,
  TimestamptzHandlingMode,
  };
  use crate::sink::formatter::AppendOnlyFormatter;
@@ -778,6 +778,7 @@ mod test {
  TimestampHandlingMode::Milli,
  TimestamptzHandlingMode::UtcString,
  TimeHandlingMode::Milli,
+ JsonbHandlingMode::String,
  ),
  )),
  )

src/connector/src/sink/mqtt.rs

@@ -31,8 +31,8 @@ use with_options::WithOptions;
 
 use super::catalog::{SinkEncode, SinkFormat, SinkFormatDesc};
 use super::encoder::{
- DateHandlingMode, JsonEncoder, ProtoEncoder, ProtoHeader, RowEncoder, SerTo, TimeHandlingMode,
- TimestampHandlingMode, TimestamptzHandlingMode,
+ DateHandlingMode, JsonEncoder, JsonbHandlingMode, ProtoEncoder, ProtoHeader, RowEncoder, SerTo,
+ TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
 };
 use super::writer::AsyncTruncateSinkWriterExt;
 use super::{DummySinkCommitCoordinator, SinkWriterParam};
@@ -221,7 +221,7 @@ impl MqttSinkWriter {
  }
 
  let timestamptz_mode = TimestamptzHandlingMode::from_options(&format_desc.options)?;
-
+ let jsonb_handling_mode = JsonbHandlingMode::from_options(&format_desc.options)?;
  let encoder = match format_desc.format {
  SinkFormat::AppendOnly => match format_desc.encode {
  SinkEncode::Json => RowEncoderWrapper::Json(JsonEncoder::new(
@@ -231,6 +231,7 @@ impl MqttSinkWriter {
  TimestampHandlingMode::Milli,
  timestamptz_mode,
  TimeHandlingMode::Milli,
+ jsonb_handling_mode,
  )),
  SinkEncode::Protobuf => {
  let (descriptor, sid) = crate::schema::protobuf::fetch_descriptor(

src/connector/src/sink/nats.rs

@@ -25,7 +25,9 @@ use tokio_retry::strategy::{jitter, ExponentialBackoff};
 use tokio_retry::Retry;
 use with_options::WithOptions;
 
-use super::encoder::{DateHandlingMode, TimeHandlingMode, TimestamptzHandlingMode};
+use super::encoder::{
+ DateHandlingMode, JsonbHandlingMode, TimeHandlingMode, TimestamptzHandlingMode,
+};
 use super::utils::chunk_to_json;
 use super::{DummySinkCommitCoordinator, SinkWriterParam};
 use crate::connector_common::NatsCommon;
@@ -145,6 +147,7 @@ impl NatsSinkWriter {
  TimestampHandlingMode::Milli,
  TimestamptzHandlingMode::UtcWithoutSuffix,
  TimeHandlingMode::Milli,
+ JsonbHandlingMode::String,
  ),
  })
  }

src/connector/src/sink/snowflake.rs

@@ -30,7 +30,8 @@ use uuid::Uuid;
 use with_options::WithOptions;
 
 use super::encoder::{
- JsonEncoder, RowEncoder, TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
+ JsonEncoder, JsonbHandlingMode, RowEncoder, TimeHandlingMode, TimestampHandlingMode,
+ TimestamptzHandlingMode,
 };
 use super::snowflake_connector::{generate_s3_file_name, SnowflakeHttpClient, SnowflakeS3Client};
 use super::writer::LogSinkerOf;
@@ -231,6 +232,7 @@ impl SnowflakeSinkWriter {
  TimestampHandlingMode::String,
  TimestamptzHandlingMode::UtcString,
  TimeHandlingMode::String,
+ JsonbHandlingMode::String,
  ),
  // initial value of `epoch` will be set to 0
  epoch: 0,