mz_interchange/
json.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::collections::{BTreeMap, BTreeSet};
11use std::fmt;
12
13use itertools::Itertools;
14use mz_repr::adt::array::ArrayDimension;
15use mz_repr::adt::char;
16use mz_repr::adt::jsonb::JsonbRef;
17use mz_repr::adt::numeric::{NUMERIC_AGG_MAX_PRECISION, NUMERIC_DATUM_MAX_PRECISION};
18use mz_repr::{CatalogItemId, ColumnName, Datum, RelationDesc, SqlColumnType, SqlScalarType};
19use serde_json::{Map, json};
20
21use crate::avro::DocTarget;
22use crate::encode::{Encode, TypedDatum, column_names_and_types};
23use crate::envelopes;
24
25const AVRO_NAMESPACE: &str = "com.materialize.sink";
26const MICROS_PER_MILLIS: u32 = 1_000;
27
28// Manages encoding of JSON-encoded bytes
29pub struct JsonEncoder {
30    columns: Vec<(ColumnName, SqlColumnType)>,
31}
32
33impl JsonEncoder {
34    pub fn new(desc: RelationDesc, debezium: bool) -> Self {
35        let mut columns = column_names_and_types(desc);
36        if debezium {
37            columns = envelopes::dbz_envelope(columns);
38        };
39        JsonEncoder { columns }
40    }
41}
42
43impl Encode for JsonEncoder {
44    fn encode_unchecked(&self, row: mz_repr::Row) -> Vec<u8> {
45        let value = encode_datums_as_json(row.iter(), self.columns.as_ref());
46        value.to_string().into_bytes()
47    }
48}
49
50impl fmt::Debug for JsonEncoder {
51    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
52        f.debug_struct("JsonEncoder")
53            .field(
54                "schema",
55                &format!(
56                    "{:?}",
57                    build_row_schema_json(
58                        &self.columns,
59                        "schema",
60                        &BTreeMap::new(),
61                        None,
62                        &Default::default(),
63                    )
64                ),
65            )
66            .finish()
67    }
68}
69
70/// Encodes a sequence of `Datum` as JSON, using supplied column names and types.
71pub fn encode_datums_as_json<'a, I>(
72    datums: I,
73    names_types: &[(ColumnName, SqlColumnType)],
74) -> serde_json::Value
75where
76    I: IntoIterator<Item = Datum<'a>>,
77{
78    let value_fields = datums
79        .into_iter()
80        .zip_eq(names_types)
81        .map(|(datum, (name, typ))| {
82            (
83                name.to_string(),
84                TypedDatum::new(datum, typ).json(&JsonNumberPolicy::KeepAsNumber),
85            )
86        })
87        .collect();
88    serde_json::Value::Object(value_fields)
89}
90
91/// Policies for how to handle Numbers in JSON.
92#[derive(Debug)]
93pub enum JsonNumberPolicy {
94    /// Do not change Numbers.
95    KeepAsNumber,
96    /// Convert Numbers to their String representation. Useful for JavaScript consumers that may
97    /// interpret some numbers incorrectly.
98    ConvertNumberToString,
99}
100
101pub trait ToJson {
102    /// Transforms this value to a JSON value.
103    fn json(self, number_policy: &JsonNumberPolicy) -> serde_json::Value;
104}
105
106impl ToJson for TypedDatum<'_> {
107    fn json(self, number_policy: &JsonNumberPolicy) -> serde_json::Value {
108        let TypedDatum { datum, typ } = self;
109        if typ.nullable && datum.is_null() {
110            return serde_json::Value::Null;
111        }
112        let value = match &typ.scalar_type {
113            SqlScalarType::AclItem => json!(datum.unwrap_acl_item().to_string()),
114            SqlScalarType::Bool => json!(datum.unwrap_bool()),
115            SqlScalarType::PgLegacyChar => json!(datum.unwrap_uint8()),
116            SqlScalarType::Int16 => json!(datum.unwrap_int16()),
117            SqlScalarType::Int32 => json!(datum.unwrap_int32()),
118            SqlScalarType::Int64 => json!(datum.unwrap_int64()),
119            SqlScalarType::UInt16 => json!(datum.unwrap_uint16()),
120            // NOTE: `regproc` stays a number here even though the pgwire text
121            // encoding resolves it to a function name. `build_row_schema_json`
122            // below declares these columns as 4-byte unsigned integers, so a
123            // name would contradict the schema this module publishes for them.
124            SqlScalarType::UInt32
125            | SqlScalarType::Oid
126            | SqlScalarType::RegClass
127            | SqlScalarType::RegProc
128            | SqlScalarType::RegType => {
129                json!(datum.unwrap_uint32())
130            }
131            SqlScalarType::UInt64 => json!(datum.unwrap_uint64()),
132            SqlScalarType::Float32 => json!(datum.unwrap_float32()),
133            SqlScalarType::Float64 => json!(datum.unwrap_float64()),
134            SqlScalarType::Numeric { .. } => {
135                json!(datum.unwrap_numeric().0.to_standard_notation_string())
136            }
137            // https://stackoverflow.com/questions/10286204/what-is-the-right-json-date-format
138            SqlScalarType::Date => serde_json::Value::String(format!("{}", datum.unwrap_date())),
139            SqlScalarType::Time => serde_json::Value::String(format!("{:?}", datum.unwrap_time())),
140            SqlScalarType::Timestamp { .. } => {
141                let dt = datum.unwrap_timestamp().to_naive().and_utc();
142                let millis = dt.timestamp_millis();
143                let micros = dt.timestamp_subsec_micros()
144                    - (dt.timestamp_subsec_millis() * MICROS_PER_MILLIS);
145                serde_json::Value::String(format!("{millis}.{micros:0>3}"))
146            }
147            SqlScalarType::TimestampTz { .. } => {
148                let dt = datum.unwrap_timestamptz().to_utc();
149                let millis = dt.timestamp_millis();
150                let micros = dt.timestamp_subsec_micros()
151                    - (dt.timestamp_subsec_millis() * MICROS_PER_MILLIS);
152                serde_json::Value::String(format!("{millis}.{micros:0>3}"))
153            }
154            SqlScalarType::Interval => {
155                serde_json::Value::String(format!("{}", datum.unwrap_interval()))
156            }
157            SqlScalarType::Bytes => json!(datum.unwrap_bytes()),
158            SqlScalarType::String | SqlScalarType::VarChar { .. } | SqlScalarType::PgLegacyName => {
159                json!(datum.unwrap_str())
160            }
161            SqlScalarType::Char { length } => {
162                let s = char::format_str_pad(datum.unwrap_str(), *length);
163                serde_json::Value::String(s)
164            }
165            SqlScalarType::Jsonb => JsonbRef::from_datum(datum).to_serde_json(),
166            SqlScalarType::Uuid => json!(datum.unwrap_uuid()),
167            ty @ (SqlScalarType::Array(..) | SqlScalarType::Int2Vector) => {
168                let array = datum.unwrap_array();
169                let dims = array.dims().into_iter().collect::<Vec<_>>();
170                let mut datums = array.elements().iter();
171                encode_array(&mut datums, &dims, &mut |datum| {
172                    TypedDatum::new(
173                        datum,
174                        &SqlColumnType {
175                            nullable: true,
176                            scalar_type: ty.unwrap_collection_element_type().clone(),
177                        },
178                    )
179                    .json(number_policy)
180                })
181            }
182            SqlScalarType::List { element_type, .. } => {
183                let values = datum
184                    .unwrap_list()
185                    .into_iter()
186                    .map(|datum| {
187                        TypedDatum::new(
188                            datum,
189                            &SqlColumnType {
190                                nullable: true,
191                                scalar_type: (**element_type).clone(),
192                            },
193                        )
194                        .json(number_policy)
195                    })
196                    .collect();
197                serde_json::Value::Array(values)
198            }
199            SqlScalarType::Record { fields, .. } => {
200                let list = datum.unwrap_list();
201                let fields: Map<String, serde_json::Value> = fields
202                    .iter()
203                    .zip_eq(list)
204                    .map(|((name, typ), datum)| {
205                        let name = name.to_string();
206                        let datum = TypedDatum::new(datum, typ);
207                        let value = datum.json(number_policy);
208                        (name, value)
209                    })
210                    .collect();
211                fields.into()
212            }
213            SqlScalarType::Map { value_type, .. } => {
214                let map = datum.unwrap_map();
215                let elements = map
216                    .into_iter()
217                    .map(|(key, datum)| {
218                        let value = TypedDatum::new(
219                            datum,
220                            &SqlColumnType {
221                                nullable: true,
222                                scalar_type: (**value_type).clone(),
223                            },
224                        )
225                        .json(number_policy);
226                        (key.to_string(), value)
227                    })
228                    .collect();
229                serde_json::Value::Object(elements)
230            }
231            SqlScalarType::MzTimestamp => json!(datum.unwrap_mz_timestamp().to_string()),
232            SqlScalarType::Range { .. } => {
233                // Ranges' interiors are not expected to be types whose
234                // string representations are misleading/wrong, e.g.
235                // records.
236                json!(datum.unwrap_range().to_string())
237            }
238            SqlScalarType::MzAclItem => json!(datum.unwrap_mz_acl_item().to_string()),
239        };
240        // We don't need to recurse into map or object here because those already recursively call
241        // .json() with the number policy to generate the member Values.
242        match (number_policy, value) {
243            (JsonNumberPolicy::KeepAsNumber, value) => value,
244            (JsonNumberPolicy::ConvertNumberToString, serde_json::Value::Number(n)) => {
245                serde_json::Value::String(n.to_string())
246            }
247            (JsonNumberPolicy::ConvertNumberToString, value) => value,
248        }
249    }
250}
251
252fn encode_array<'a>(
253    elems: &mut impl Iterator<Item = Datum<'a>>,
254    dims: &[ArrayDimension],
255    elem_encoder: &mut impl FnMut(Datum<'_>) -> serde_json::Value,
256) -> serde_json::Value {
257    serde_json::Value::Array(match dims {
258        [] => vec![],
259        [dim] => elems.take(dim.length).map(elem_encoder).collect(),
260        [dim, rest @ ..] => (0..dim.length)
261            .map(|_| encode_array(elems, rest, elem_encoder))
262            .collect(),
263    })
264}
265
266fn build_row_schema_field_type(
267    type_namer: &mut Namer,
268    custom_names: &BTreeMap<CatalogItemId, String>,
269    typ: &SqlColumnType,
270    item_id: Option<CatalogItemId>,
271    options: &SchemaOptions,
272) -> serde_json::Value {
273    let mut field_type = match &typ.scalar_type {
274        SqlScalarType::AclItem => json!("string"),
275        SqlScalarType::Bool => json!("boolean"),
276        SqlScalarType::PgLegacyChar => json!({
277            "type": "fixed",
278            "size": 1,
279        }),
280        SqlScalarType::Int16 | SqlScalarType::Int32 => {
281            json!("int")
282        }
283        SqlScalarType::Int64 => json!("long"),
284        SqlScalarType::UInt16 => type_namer.unsigned_type(2),
285        SqlScalarType::UInt32
286        | SqlScalarType::Oid
287        | SqlScalarType::RegClass
288        | SqlScalarType::RegProc
289        | SqlScalarType::RegType => type_namer.unsigned_type(4),
290        SqlScalarType::UInt64 => type_namer.unsigned_type(8),
291        SqlScalarType::Float32 => json!("float"),
292        SqlScalarType::Float64 => json!("double"),
293        SqlScalarType::Date => json!({
294            "type": "int",
295            "logicalType": "date",
296        }),
297        SqlScalarType::Time => json!({
298            "type": "long",
299            "logicalType": "time-micros",
300        }),
301        SqlScalarType::Timestamp { precision } | SqlScalarType::TimestampTz { precision } => {
302            json!({
303                "type": "long",
304                "logicalType": match precision {
305                    Some(precision) if precision.into_u8() <= 3 => "timestamp-millis",
306                    _ => "timestamp-micros",
307                },
308            })
309        }
310        SqlScalarType::Interval => type_namer.interval_type(),
311        SqlScalarType::Bytes => json!("bytes"),
312        SqlScalarType::String
313        | SqlScalarType::Char { .. }
314        | SqlScalarType::VarChar { .. }
315        | SqlScalarType::PgLegacyName => {
316            json!("string")
317        }
318        SqlScalarType::Jsonb => json!({
319            "type": "string",
320            "connect.name": "io.debezium.data.Json",
321        }),
322        SqlScalarType::Uuid => json!({
323            "type": "string",
324            "logicalType": "uuid",
325        }),
326        ty
327        @ (SqlScalarType::Array(..) | SqlScalarType::Int2Vector | SqlScalarType::List { .. }) => {
328            let inner = build_row_schema_field_type(
329                type_namer,
330                custom_names,
331                &SqlColumnType {
332                    nullable: true,
333                    scalar_type: ty.unwrap_collection_element_type().clone(),
334                },
335                item_id,
336                options,
337            );
338            json!({
339                "type": "array",
340                "items": inner
341            })
342        }
343        SqlScalarType::Map { value_type, .. } => {
344            let inner = build_row_schema_field_type(
345                type_namer,
346                custom_names,
347                &SqlColumnType {
348                    nullable: true,
349                    scalar_type: (**value_type).clone(),
350                },
351                item_id,
352                options,
353            );
354            json!({
355                "type": "map",
356                "values": inner
357            })
358        }
359        SqlScalarType::Record {
360            fields, custom_id, ..
361        } => {
362            let (name, name_seen) = match custom_id.as_ref().and_then(|id| custom_names.get(id)) {
363                Some(name) => type_namer.valid_name(name),
364                None => (type_namer.anonymous_record_name(), false),
365            };
366            if name_seen {
367                json!(name)
368            } else {
369                let fields = fields.to_vec();
370                let json_fields =
371                    build_row_schema_fields(&fields, type_namer, custom_names, *custom_id, options);
372                if let Some(comment) =
373                    custom_id.and_then(|id| options.doc_comments.get(&DocTarget::Type(id)))
374                {
375                    json!({
376                        "type": "record",
377                        "name": name,
378                        "doc": comment,
379                        "fields": json_fields
380                    })
381                } else {
382                    json!({
383                        "type": "record",
384                        "name": name,
385                        "fields": json_fields
386                    })
387                }
388            }
389        }
390        SqlScalarType::Numeric { max_scale } => {
391            let (p, s) = match max_scale {
392                Some(max_scale) => (NUMERIC_DATUM_MAX_PRECISION, max_scale.into_u8()),
393                None => (NUMERIC_AGG_MAX_PRECISION, NUMERIC_DATUM_MAX_PRECISION),
394            };
395            json!({
396                "type": "bytes",
397                "logicalType": "decimal",
398                "precision": p,
399                "scale": s,
400            })
401        }
402        SqlScalarType::MzTimestamp => json!("string"),
403        // https://debezium.io/documentation/reference/stable/connectors/postgresql.html
404        SqlScalarType::Range { .. } => json!("string"),
405        SqlScalarType::MzAclItem => json!("string"),
406    };
407    if typ.nullable {
408        // Should be revisited if we ever support a different kind of union scheme.
409        // Currently adding the "null" at the beginning means we can set the default
410        // value to "null" if such a preference is set.
411        field_type = json!(["null", field_type]);
412    }
413    field_type
414}
415
416fn build_row_schema_fields(
417    columns: &[(ColumnName, SqlColumnType)],
418    type_namer: &mut Namer,
419    custom_names: &BTreeMap<CatalogItemId, String>,
420    item_id: Option<CatalogItemId>,
421    options: &SchemaOptions,
422) -> Vec<serde_json::Value> {
423    let mut fields = Vec::new();
424    let mut field_namer = Namer::default();
425    for (name, typ) in columns.iter() {
426        let (name, _seen) = field_namer.valid_name(name);
427        let field_type =
428            build_row_schema_field_type(type_namer, custom_names, typ, item_id, options);
429
430        let mut field = json!({
431            "name": name,
432            "type": field_type,
433        });
434
435        // It's a nullable union if the type is an array and the first option is "null"
436        let is_nullable_union = field_type
437            .as_array()
438            .is_some_and(|array| array.first().is_some_and(|first| first == &json!("null")));
439
440        if options.set_null_defaults && is_nullable_union {
441            field
442                .as_object_mut()
443                .expect("`field` initialized to JSON object above")
444                .insert("default".to_string(), json!(null));
445        }
446
447        if let Some(comment) = item_id.and_then(|item_id| {
448            options.doc_comments.get(&DocTarget::Field {
449                object_id: item_id,
450                column_name: name.into(),
451            })
452        }) {
453            field
454                .as_object_mut()
455                .expect("`field` initialized to JSON object above")
456                .insert("doc".to_string(), json!(comment));
457        }
458
459        fields.push(field);
460    }
461    fields
462}
463
464#[derive(Default, Clone, Debug)]
465/// Struct to pass around options to create the json schema
466pub struct SchemaOptions {
467    /// Boolean flag to enable null defaults.
468    pub set_null_defaults: bool,
469    /// Map containing comments for an item or field, used to populate
470    /// documentation in the generated avro schema
471    pub doc_comments: BTreeMap<DocTarget, String>,
472}
473
474/// Builds the JSON for the row schema, which can be independently useful.
475pub fn build_row_schema_json(
476    columns: &[(ColumnName, SqlColumnType)],
477    name: &str,
478    custom_names: &BTreeMap<CatalogItemId, String>,
479    item_id: Option<CatalogItemId>,
480    options: &SchemaOptions,
481) -> Result<serde_json::Value, anyhow::Error> {
482    let fields = build_row_schema_fields(
483        columns,
484        &mut Namer::default(),
485        custom_names,
486        item_id,
487        options,
488    );
489
490    let _ = mz_avro::schema::Name::parse_simple(name)?;
491    if let Some(comment) =
492        item_id.and_then(|item_id| options.doc_comments.get(&DocTarget::Type(item_id)))
493    {
494        Ok(json!({
495            "type": "record",
496            "doc": comment,
497            "fields": fields,
498            "name": name
499        }))
500    } else {
501        Ok(json!({
502            "type": "record",
503            "fields": fields,
504            "name": name
505        }))
506    }
507}
508
509/// Naming helper for use when constructing an Avro schema.
510#[derive(Default)]
511struct Namer {
512    record_index: usize,
513    seen_interval: bool,
514    seen_unsigneds: BTreeSet<usize>,
515    seen_names: BTreeMap<String, String>,
516    valid_names_count: BTreeMap<String, usize>,
517}
518
519impl Namer {
520    /// Returns the schema for an interval type.
521    fn interval_type(&mut self) -> serde_json::Value {
522        let name = format!("{AVRO_NAMESPACE}.interval");
523        if self.seen_interval {
524            json!(name)
525        } else {
526            self.seen_interval = true;
527            json!({
528            "type": "fixed",
529            "size": 16,
530            "name": name,
531            })
532        }
533    }
534
535    /// Returns the schema for an unsigned integer with the given width.
536    fn unsigned_type(&mut self, width: usize) -> serde_json::Value {
537        let name = format!("{AVRO_NAMESPACE}.uint{width}");
538        if self.seen_unsigneds.contains(&width) {
539            json!(name)
540        } else {
541            self.seen_unsigneds.insert(width);
542            json!({
543                "type": "fixed",
544                "size": width,
545                "name": name,
546            })
547        }
548    }
549
550    /// Returns a name to use for a new anonymous record.
551    fn anonymous_record_name(&mut self) -> String {
552        let out = format!("{AVRO_NAMESPACE}.record{}", self.record_index);
553        self.record_index += 1;
554        out
555    }
556
557    /// Turns `name` into a valid, unique name for use in the Avro schema.
558    ///
559    /// Returns the valid name and whether `name` has been seen before.
560    fn valid_name(&mut self, name: &str) -> (String, bool) {
561        if let Some(valid_name) = self.seen_names.get(name) {
562            (valid_name.into(), true)
563        } else {
564            let mut valid_name = mz_avro::schema::Name::make_valid(name);
565            let valid_name_count = self
566                .valid_names_count
567                .entry(valid_name.clone())
568                .or_default();
569            if *valid_name_count != 0 {
570                valid_name += &valid_name_count.to_string();
571            }
572            *valid_name_count += 1;
573            self.seen_names.insert(name.into(), valid_name.clone());
574            (valid_name, false)
575        }
576    }
577}
mz_interchange/json.rs

mz_interchange/
json.rs
