risingwave_connector/parser/
json_parser.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// Note on this file:
16//
17// There's no struct named `JsonParser` anymore since #13707. `ENCODE JSON` will be
18// dispatched to `PlainParser` or `UpsertParser` with `JsonAccessBuilder` instead.
19//
20// This file now only contains utilities and tests for JSON parsing. Also, to avoid
21// rely on the internal implementation and allow that to be changed, the tests use
22// `ByteStreamSourceParserImpl` to create a parser instance.
23
24use std::collections::BTreeMap;
25
26use anyhow::Context as _;
27use risingwave_common::bail;
28use risingwave_common::catalog::Field;
29use risingwave_connector_codec::JsonSchema;
30use risingwave_connector_codec::decoder::avro::avro_schema_to_fields;
31
32use super::utils::{bytes_from_url, get_kafka_topic, get_pulsar_topic};
33use super::{JsonProperties, SchemaRegistryConfig};
34use crate::error::ConnectorResult;
35use crate::parser::AccessBuilder;
36use crate::parser::unified::AccessImpl;
37use crate::parser::unified::json::{JsonAccess, JsonParseOptions};
38use crate::schema::pulsar_schema::{
39    self, PulsarSchema, PulsarSchemaConfig, parse_partitioned_topic,
40};
41use crate::schema::schema_registry::{Client, handle_sr_list};
42
43#[derive(Debug)]
44pub struct JsonAccessBuilder {
45    value: Option<Vec<u8>>,
46    payload_start_idx: usize,
47    json_parse_options: JsonParseOptions,
48}
49
50impl JsonAccessBuilder {
51    pub fn generate_json_access(&mut self, payload: Vec<u8>) -> ConnectorResult<JsonAccess<'_>> {
52        // XXX: When will we enter this branch?
53        if payload.is_empty() {
54            self.value = Some("{}".into());
55        } else {
56            self.value = Some(payload);
57        }
58        let value = simd_json::to_borrowed_value(
59            &mut self.value.as_mut().unwrap()[self.payload_start_idx..],
60        )
61        .context("failed to parse json payload")?;
62        Ok(JsonAccess::new_with_options(
63            value,
64            // Debezium and Canal have their special json access builder and will not
65            // use this
66            &self.json_parse_options,
67        ))
68    }
69}
70
71impl AccessBuilder for JsonAccessBuilder {
72    async fn generate_accessor(
73        &mut self,
74        payload: Vec<u8>,
75        _: &crate::source::SourceMeta,
76    ) -> ConnectorResult<AccessImpl<'_>> {
77        Ok(AccessImpl::Json(self.generate_json_access(payload)?))
78    }
79}
80
81impl JsonAccessBuilder {
82    pub fn new(config: JsonProperties) -> ConnectorResult<Self> {
83        let mut json_parse_options = JsonParseOptions::DEFAULT;
84        if let Some(mode) = config.timestamp_handling {
85            json_parse_options.timestamp_handling = mode;
86        }
87        if let Some(mode) = config.timestamptz_handling {
88            json_parse_options.timestamptz_handling = mode;
89        }
90        if let Some(mode) = config.time_handling {
91            json_parse_options.time_handling = mode;
92        }
93        if let Some(mode) = config.bigint_unsigned_handling {
94            json_parse_options.bigint_unsigned_handling = mode;
95        }
96        json_parse_options.handle_toast_columns = config.handle_toast_columns;
97        Ok(Self {
98            value: None,
99            payload_start_idx: if config.use_schema_registry { 5 } else { 0 },
100            json_parse_options,
101        })
102    }
103}
104
105pub async fn fetch_json_schema_and_map_to_columns(
106    schema_location: &str,
107    schema_registry_auth: Option<SchemaRegistryConfig>,
108    props: &BTreeMap<String, String>,
109) -> ConnectorResult<Vec<Field>> {
110    let url = handle_sr_list(schema_location)?;
111    let mut json_schema = if let Some(schema_registry_auth) = schema_registry_auth {
112        let client = Client::new(url.clone(), &schema_registry_auth)?;
113        let topic = get_kafka_topic(props)?;
114        let schema = client
115            .get_schema_by_subject(&format!("{}-value", topic))
116            .await?;
117        JsonSchema::parse_str(&schema.content)?
118    } else {
119        let url = url.first().unwrap();
120        let bytes = bytes_from_url(url, None).await?;
121        JsonSchema::parse_bytes(&bytes)?
122    };
123    json_schema
124        .json_schema_to_columns(url.first().unwrap().clone())
125        .await
126        .map_err(Into::into)
127}
128
129/// Fetches the latest Pulsar `JSON` schema of the source topic and maps it to columns.
130///
131/// Pulsar stores `JSON` schemas as Avro record definitions, while payloads are plain JSON without
132/// a schema-registry header. JSON is self-describing, so only the latest schema is needed.
133pub async fn fetch_pulsar_json_schema_and_map_to_columns(
134    config: &PulsarSchemaConfig,
135    props: &BTreeMap<String, String>,
136) -> ConnectorResult<Vec<Field>> {
137    let client = pulsar_schema::Client::new(config)?;
138    let topic = parse_partitioned_topic(get_pulsar_topic(props)?)?;
139    match client.get_schema(&topic, None).await?.parse()? {
140        PulsarSchema::Json(schema) => Ok(avro_schema_to_fields(&schema, None)?),
141        PulsarSchema::Avro(_) => bail!("expected Pulsar JSON schema, got AVRO"),
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use std::vec;
148
149    use itertools::Itertools;
150    use risingwave_common::array::{Op, StructValue};
151    use risingwave_common::catalog::ColumnDesc;
152    use risingwave_common::row::Row;
153    use risingwave_common::test_prelude::StreamChunkTestExt;
154    use risingwave_common::types::{DataType, ScalarImpl, StructType, ToOwnedDatum};
155    use risingwave_pb::plan_common::additional_column::ColumnType as AdditionalColumnType;
156    use risingwave_pb::plan_common::{AdditionalColumn, AdditionalColumnKey};
157
158    use crate::parser::test_utils::ByteStreamSourceParserImplTestExt as _;
159    use crate::parser::{
160        ByteStreamSourceParserImpl, CommonParserConfig, ParserConfig, ProtocolProperties,
161        SourceColumnDesc, SpecificParserConfig,
162    };
163    use crate::source::SourceColumnType;
164
165    fn make_parser(rw_columns: Vec<SourceColumnDesc>) -> ByteStreamSourceParserImpl {
166        ByteStreamSourceParserImpl::create_for_test(ParserConfig {
167            common: CommonParserConfig { rw_columns },
168            specific: SpecificParserConfig::DEFAULT_PLAIN_JSON,
169        })
170        .unwrap()
171    }
172
173    fn make_upsert_parser(rw_columns: Vec<SourceColumnDesc>) -> ByteStreamSourceParserImpl {
174        ByteStreamSourceParserImpl::create_for_test(ParserConfig {
175            common: CommonParserConfig { rw_columns },
176            specific: SpecificParserConfig {
177                protocol_config: ProtocolProperties::Upsert,
178                ..SpecificParserConfig::DEFAULT_PLAIN_JSON
179            },
180        })
181        .unwrap()
182    }
183
184    fn get_payload() -> Vec<Vec<u8>> {
185        vec![
186            br#"{"i32":1,"bool":true,"i16":1,"i64":12345678,"f32":1.23,"f64":1.2345,"varchar":"varchar","date":"2021-01-01","timestamp":"2021-01-01 16:06:12.269","decimal":12345.67890,"interval":"P1Y2M3DT0H5M0S"}"#.to_vec(),
187            br#"{"i32":1,"f32":12345e+10,"f64":12345,"decimal":12345,"interval":"1 day"}"#.to_vec(),
188        ]
189    }
190
191    fn get_array_top_level_payload() -> Vec<Vec<u8>> {
192        vec![
193            br#"[{"i32":1,"bool":true,"i16":1,"i64":12345678,"f32":1.23,"f64":1.2345,"varchar":"varchar","date":"2021-01-01","timestamp":"2021-01-01 16:06:12.269","decimal":12345.67890}, {"i32":1,"f32":12345e+10,"f64":12345,"decimal":12345}]"#.to_vec()
194        ]
195    }
196
197    async fn test_json_parser(get_payload: fn() -> Vec<Vec<u8>>) {
198        let descs = vec![
199            SourceColumnDesc::simple("i32", DataType::Int32, 0.into()),
200            SourceColumnDesc::simple("bool", DataType::Boolean, 2.into()),
201            SourceColumnDesc::simple("i16", DataType::Int16, 3.into()),
202            SourceColumnDesc::simple("i64", DataType::Int64, 4.into()),
203            SourceColumnDesc::simple("f32", DataType::Float32, 5.into()),
204            SourceColumnDesc::simple("f64", DataType::Float64, 6.into()),
205            SourceColumnDesc::simple("varchar", DataType::Varchar, 7.into()),
206            SourceColumnDesc::simple("date", DataType::Date, 8.into()),
207            SourceColumnDesc::simple("timestamp", DataType::Timestamp, 9.into()),
208            SourceColumnDesc::simple("decimal", DataType::Decimal, 10.into()),
209            SourceColumnDesc::simple("interval", DataType::Interval, 11.into()),
210        ];
211
212        let parser = make_parser(descs);
213        let chunk = parser.parse(get_payload()).await;
214
215        let mut rows = chunk.rows();
216
217        {
218            let (op, row) = rows.next().unwrap();
219            assert_eq!(op, Op::Insert);
220            assert_eq!(row.datum_at(0).to_owned_datum(), Some(ScalarImpl::Int32(1)));
221            assert_eq!(
222                row.datum_at(1).to_owned_datum(),
223                (Some(ScalarImpl::Bool(true)))
224            );
225            assert_eq!(
226                row.datum_at(2).to_owned_datum(),
227                (Some(ScalarImpl::Int16(1)))
228            );
229            assert_eq!(
230                row.datum_at(3).to_owned_datum(),
231                (Some(ScalarImpl::Int64(12345678)))
232            );
233            assert_eq!(
234                row.datum_at(4).to_owned_datum(),
235                (Some(ScalarImpl::Float32(1.23.into())))
236            );
237            assert_eq!(
238                row.datum_at(5).to_owned_datum(),
239                (Some(ScalarImpl::Float64(1.2345.into())))
240            );
241            assert_eq!(
242                row.datum_at(6).to_owned_datum(),
243                (Some(ScalarImpl::Utf8("varchar".into())))
244            );
245            assert_eq!(
246                row.datum_at(7).to_owned_datum(),
247                (Some(ScalarImpl::Date("2021-01-01".parse().unwrap())))
248            );
249            assert_eq!(
250                row.datum_at(8).to_owned_datum(),
251                (Some(ScalarImpl::Timestamp(
252                    "2021-01-01 16:06:12.269".parse().unwrap()
253                )))
254            );
255            assert_eq!(
256                row.datum_at(9).to_owned_datum(),
257                (Some(ScalarImpl::Decimal("12345.67890".parse().unwrap())))
258            );
259            assert_eq!(
260                row.datum_at(10).to_owned_datum(),
261                (Some(ScalarImpl::Interval("P1Y2M3DT0H5M0S".parse().unwrap())))
262            );
263        }
264
265        {
266            let (op, row) = rows.next().unwrap();
267            assert_eq!(op, Op::Insert);
268            assert_eq!(
269                row.datum_at(0).to_owned_datum(),
270                (Some(ScalarImpl::Int32(1)))
271            );
272            assert_eq!(row.datum_at(1).to_owned_datum(), None);
273            assert_eq!(
274                row.datum_at(4).to_owned_datum(),
275                (Some(ScalarImpl::Float32(12345e+10.into())))
276            );
277            assert_eq!(
278                row.datum_at(5).to_owned_datum(),
279                (Some(ScalarImpl::Float64(12345.into())))
280            );
281            assert_eq!(
282                row.datum_at(9).to_owned_datum(),
283                (Some(ScalarImpl::Decimal(12345.into())))
284            );
285            assert_eq!(
286                row.datum_at(10).to_owned_datum(),
287                (Some(ScalarImpl::Interval("1 day".parse().unwrap())))
288            );
289        }
290    }
291
292    #[tokio::test]
293    async fn test_json_parse_object_top_level() {
294        test_json_parser(get_payload).await;
295    }
296    #[ignore]
297    #[tokio::test]
298    async fn test_json_parse_array_top_level() {
299        test_json_parser(get_array_top_level_payload).await;
300    }
301
302    #[tokio::test]
303    async fn test_json_parser_failed() {
304        let descs = vec![
305            SourceColumnDesc::simple("v1", DataType::Int32, 0.into()),
306            SourceColumnDesc::simple("v2", DataType::Int16, 1.into()),
307            SourceColumnDesc::simple("v3", DataType::Varchar, 2.into()),
308        ];
309
310        let parser = make_parser(descs);
311        let payloads = vec![
312            // Parse a correct record.
313            br#"{"v1": 1, "v2": 2, "v3": "3"}"#.to_vec(),
314            // Parse an incorrect record.
315            // `v2` overflowed.
316            // ignored the error, and fill None at v2.
317            br#"{"v1": 1, "v2": 65536, "v3": "3"}"#.to_vec(),
318            // Parse a correct record.
319            br#"{"v1": 1, "v2": 2, "v3": "3"}"#.to_vec(),
320        ];
321        let chunk = parser.parse(payloads).await;
322
323        assert!(chunk.valid());
324        assert_eq!(chunk.cardinality(), 3);
325
326        let row_vec = chunk.rows().collect_vec();
327        assert_eq!(row_vec[1].1.datum_at(1), None);
328    }
329
330    #[tokio::test]
331    async fn test_json_parse_struct() {
332        let descs = [
333            ColumnDesc::named(
334                "data",
335                0.into(),
336                DataType::from(StructType::new([
337                    ("created_at", DataType::Timestamp),
338                    ("id", DataType::Varchar),
339                    ("text", DataType::Varchar),
340                    ("lang", DataType::Varchar),
341                ])),
342            ),
343            ColumnDesc::named(
344                "author",
345                5.into(),
346                DataType::from(StructType::new([
347                    ("created_at", DataType::Timestamp),
348                    ("id", DataType::Varchar),
349                    ("name", DataType::Varchar),
350                    ("username", DataType::Varchar),
351                ])),
352            ),
353            ColumnDesc::named("I64CastToVarchar", 10.into(), DataType::Varchar),
354            ColumnDesc::named("VarcharCastToI64", 11.into(), DataType::Int64),
355        ]
356        .iter()
357        .map(SourceColumnDesc::from)
358        .collect_vec();
359
360        let parser = make_parser(descs);
361        let payload = br#"
362        {
363            "data": {
364                "created_at": "2022-07-13 20:48:37.07",
365                "id": "1732524418112319151",
366                "text": "Here man favor ourselves mysteriously most her sigh in straightaway for afterwards.",
367                "lang": "English"
368            },
369            "author": {
370                "created_at": "2018-01-29 12:19:11.07",
371                "id": "7772634297",
372                "name": "Lily Frami yet",
373                "username": "Dooley5659"
374            },
375            "I64CastToVarchar": 1598197865760800768,
376            "VarcharCastToI64": "1598197865760800768"
377        }
378        "#.to_vec();
379        let chunk = parser.parse(vec![payload]).await;
380
381        let (op, row) = chunk.rows().next().unwrap();
382        assert_eq!(op, Op::Insert);
383        let row = row.into_owned_row().into_inner();
384
385        let expected = vec![
386            Some(ScalarImpl::Struct(StructValue::new(vec![
387                Some(ScalarImpl::Timestamp(
388                    "2022-07-13 20:48:37.07".parse().unwrap()
389                )),
390                Some(ScalarImpl::Utf8("1732524418112319151".into())),
391                Some(ScalarImpl::Utf8("Here man favor ourselves mysteriously most her sigh in straightaway for afterwards.".into())),
392                Some(ScalarImpl::Utf8("English".into())),
393            ]))),
394            Some(ScalarImpl::Struct(StructValue::new(vec![
395                Some(ScalarImpl::Timestamp(
396                    "2018-01-29 12:19:11.07".parse().unwrap()
397                )),
398                Some(ScalarImpl::Utf8("7772634297".into())),
399                Some(ScalarImpl::Utf8("Lily Frami yet".into())),
400                Some(ScalarImpl::Utf8("Dooley5659".into())),
401            ]) )),
402            Some(ScalarImpl::Utf8("1598197865760800768".into())),
403            Some(ScalarImpl::Int64(1598197865760800768)),
404        ];
405        assert_eq!(row, expected.into());
406    }
407
408    #[tokio::test]
409    async fn test_json_parse_struct_from_string() {
410        let descs = [ColumnDesc::named(
411            "struct",
412            0.into(),
413            DataType::from(StructType::new([
414                ("varchar", DataType::Varchar),
415                ("boolean", DataType::Boolean),
416            ])),
417        )]
418        .iter()
419        .map(SourceColumnDesc::from)
420        .collect_vec();
421
422        let parser = make_parser(descs);
423        let payload = br#"
424        {
425            "struct": "{\"varchar\": \"varchar\", \"boolean\": true}"
426        }
427        "#
428        .to_vec();
429        let chunk = parser.parse(vec![payload]).await;
430
431        let (op, row) = chunk.rows().next().unwrap();
432        assert_eq!(op, Op::Insert);
433        let row = row.into_owned_row().into_inner();
434
435        let expected = vec![Some(ScalarImpl::Struct(StructValue::new(vec![
436            Some(ScalarImpl::Utf8("varchar".into())),
437            Some(ScalarImpl::Bool(true)),
438        ])))];
439        assert_eq!(row, expected.into());
440    }
441
442    #[cfg(not(madsim))] // Traced test does not work with madsim
443    #[tokio::test]
444    #[tracing_test::traced_test]
445    async fn test_json_parse_struct_missing_field_warning() {
446        let descs = [ColumnDesc::named(
447            "struct",
448            0.into(),
449            DataType::from(StructType::new([
450                ("varchar", DataType::Varchar),
451                ("boolean", DataType::Boolean),
452            ])),
453        )]
454        .iter()
455        .map(SourceColumnDesc::from)
456        .collect_vec();
457
458        let parser = make_parser(descs);
459        let payload = br#"
460        {
461            "struct": {
462                "varchar": "varchar"
463            }
464        }
465        "#
466        .to_vec();
467        let chunk = parser.parse(vec![payload]).await;
468
469        let (op, row) = chunk.rows().next().unwrap();
470        assert_eq!(op, Op::Insert);
471        let row = row.into_owned_row().into_inner();
472
473        let expected = vec![Some(ScalarImpl::Struct(StructValue::new(vec![
474            Some(ScalarImpl::Utf8("varchar".into())),
475            None,
476        ])))];
477        assert_eq!(row, expected.into());
478
479        assert!(logs_contain("undefined nested field, padding with `NULL`"));
480    }
481
482    #[tokio::test]
483    async fn test_json_upsert_parser() {
484        let items = [
485            (r#"{"a":1}"#, r#"{"a":1,"b":2}"#),
486            (r#"{"a":1}"#, r#"{"a":1,"b":3}"#),
487            (r#"{"a":2}"#, r#"{"a":2,"b":2}"#),
488            (r#"{"a":2}"#, r#""#),
489        ]
490        .into_iter()
491        .map(|(k, v)| (k.as_bytes().to_vec(), v.as_bytes().to_vec()))
492        .collect_vec();
493
494        let key_column_desc = SourceColumnDesc {
495            name: "rw_key".into(),
496            data_type: DataType::Bytea,
497            column_id: 2.into(),
498            column_type: SourceColumnType::Normal,
499            is_pk: true,
500            is_hidden_addition_col: false,
501            additional_column: AdditionalColumn {
502                column_type: Some(AdditionalColumnType::Key(AdditionalColumnKey {})),
503            },
504        };
505        let descs = vec![
506            SourceColumnDesc::simple("a", DataType::Int32, 0.into()),
507            SourceColumnDesc::simple("b", DataType::Int32, 1.into()),
508            key_column_desc,
509        ];
510
511        let parser = make_upsert_parser(descs);
512        let chunk = parser.parse_upsert(items).await;
513
514        // expected chunk
515        // +---+---+---+------------------+
516        // | + | 1 | 2 | \x7b2261223a317d |
517        // | + | 1 | 3 | \x7b2261223a317d |
518        // | + | 2 | 2 | \x7b2261223a327d |
519        // | - |   |   | \x7b2261223a327d |
520        // +---+---+---+------------------+
521
522        let mut rows = chunk.rows();
523        {
524            let (op, row) = rows.next().unwrap();
525            assert_eq!(op, Op::Insert);
526            assert_eq!(
527                row.datum_at(0).to_owned_datum(),
528                (Some(ScalarImpl::Int32(1)))
529            );
530        }
531
532        {
533            let (op, row) = rows.next().unwrap();
534            assert_eq!(op, Op::Insert);
535            assert_eq!(
536                row.datum_at(0).to_owned_datum(),
537                (Some(ScalarImpl::Int32(1)))
538            );
539        }
540        {
541            let (op, row) = rows.next().unwrap();
542            assert_eq!(op, Op::Insert);
543            assert_eq!(
544                row.datum_at(0).to_owned_datum(),
545                (Some(ScalarImpl::Int32(2)))
546            );
547        }
548        {
549            let (op, row) = rows.next().unwrap();
550            assert_eq!(op, Op::Delete);
551            assert_eq!(row.datum_at(0).to_owned_datum(), (None));
552        }
553    }
554}
risingwave_connector/parser/json_parser.rs

risingwave_connector/parser/
json_parser.rs
