mz_expr/scalar/func/impls/
jsonb.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;
11use std::fmt;
12
13use mz_expr_derive::sqlfunc;
14use mz_lowertest::MzReflect;
15use mz_repr::adt::jsonb::{Jsonb, JsonbRef};
16use mz_repr::adt::mz_acl_item::{AclMode, MzAclItem};
17use mz_repr::adt::numeric::{self, Numeric, NumericMaxScale};
18use mz_repr::role_id::RoleId;
19use mz_repr::{ArrayRustType, Datum, Row, RowPacker, SqlColumnType, SqlScalarType, strconv};
20use mz_sql_parser::ast::display::AstDisplay;
21use mz_sql_parser::ast::{AstInfo, Format, FormatSpecifier, RawClusterName, RawItemName};
22use serde::{Deserialize, Serialize};
23use serde_json::json;
24
25use crate::EvalError;
26use crate::scalar::func::EagerUnaryFunc;
27use crate::scalar::func::impls::numeric::*;
28
29#[sqlfunc(
30    sqlname = "jsonb_to_text",
31    preserves_uniqueness = false,
32    inverse = to_unary!(super::CastStringToJsonb)
33)]
34pub fn cast_jsonb_to_string<'a>(a: JsonbRef<'a>) -> String {
35    let mut buf = String::new();
36    strconv::format_jsonb(&mut buf, a);
37    buf
38}
39
40#[sqlfunc(sqlname = "jsonb_to_smallint", is_monotone = true)]
41fn cast_jsonb_to_int16<'a>(a: JsonbRef<'a>) -> Result<i16, EvalError> {
42    match a.into_datum() {
43        Datum::Numeric(a) => cast_numeric_to_int16(a.into_inner()),
44        datum => Err(EvalError::InvalidJsonbCast {
45            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
46            to: "smallint".into(),
47        }),
48    }
49}
50
51#[sqlfunc(sqlname = "jsonb_to_integer", is_monotone = true)]
52fn cast_jsonb_to_int32<'a>(a: JsonbRef<'a>) -> Result<i32, EvalError> {
53    match a.into_datum() {
54        Datum::Numeric(a) => cast_numeric_to_int32(a.into_inner()),
55        datum => Err(EvalError::InvalidJsonbCast {
56            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
57            to: "integer".into(),
58        }),
59    }
60}
61
62#[sqlfunc(sqlname = "jsonb_to_bigint", is_monotone = true)]
63fn cast_jsonb_to_int64<'a>(a: JsonbRef<'a>) -> Result<i64, EvalError> {
64    match a.into_datum() {
65        Datum::Numeric(a) => cast_numeric_to_int64(a.into_inner()),
66        datum => Err(EvalError::InvalidJsonbCast {
67            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
68            to: "bigint".into(),
69        }),
70    }
71}
72
73#[sqlfunc(sqlname = "jsonb_to_real", is_monotone = true)]
74fn cast_jsonb_to_float32<'a>(a: JsonbRef<'a>) -> Result<f32, EvalError> {
75    match a.into_datum() {
76        Datum::Numeric(a) => cast_numeric_to_float32(a.into_inner()),
77        datum => Err(EvalError::InvalidJsonbCast {
78            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
79            to: "real".into(),
80        }),
81    }
82}
83
84#[sqlfunc(sqlname = "jsonb_to_double", is_monotone = true)]
85fn cast_jsonb_to_float64<'a>(a: JsonbRef<'a>) -> Result<f64, EvalError> {
86    match a.into_datum() {
87        Datum::Numeric(a) => cast_numeric_to_float64(a.into_inner()),
88        datum => Err(EvalError::InvalidJsonbCast {
89            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
90            to: "double precision".into(),
91        }),
92    }
93}
94
95#[derive(
96    Ord,
97    PartialOrd,
98    Clone,
99    Debug,
100    Eq,
101    PartialEq,
102    Serialize,
103    Deserialize,
104    Hash,
105    MzReflect
106)]
107pub struct CastJsonbToNumeric(pub Option<NumericMaxScale>);
108
109impl EagerUnaryFunc for CastJsonbToNumeric {
110    type Input<'a> = JsonbRef<'a>;
111    type Output<'a> = Result<Numeric, EvalError>;
112
113    fn call<'a>(&self, a: Self::Input<'a>) -> Self::Output<'a> {
114        match a.into_datum() {
115            Datum::Numeric(mut num) => match self.0 {
116                None => Ok(num.into_inner()),
117                Some(scale) => {
118                    if numeric::rescale(&mut num.0, scale.into_u8()).is_err() {
119                        return Err(EvalError::NumericFieldOverflow);
120                    };
121                    Ok(num.into_inner())
122                }
123            },
124            datum => Err(EvalError::InvalidJsonbCast {
125                from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
126                to: "numeric".into(),
127            }),
128        }
129    }
130
131    fn output_sql_type(&self, input: SqlColumnType) -> SqlColumnType {
132        SqlScalarType::Numeric { max_scale: self.0 }.nullable(input.nullable)
133    }
134
135    fn is_monotone(&self) -> bool {
136        true
137    }
138}
139
140impl fmt::Display for CastJsonbToNumeric {
141    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
142        f.write_str("jsonb_to_numeric")
143    }
144}
145
146#[sqlfunc(sqlname = "jsonb_to_boolean", is_monotone = true)]
147fn cast_jsonb_to_bool<'a>(a: JsonbRef<'a>) -> Result<bool, EvalError> {
148    match a.into_datum() {
149        Datum::True => Ok(true),
150        Datum::False => Ok(false),
151        datum => Err(EvalError::InvalidJsonbCast {
152            from: jsonb_typeof(JsonbRef::from_datum(datum)).into(),
153            to: "boolean".into(),
154        }),
155    }
156}
157
158#[sqlfunc(sqlname = "jsonbable_to_jsonb")]
159fn cast_jsonbable_to_jsonb<'a>(a: JsonbRef<'a>) -> JsonbRef<'a> {
160    match a.into_datum() {
161        Datum::Numeric(n) => {
162            let n = n.into_inner();
163            let datum = if n.is_finite() {
164                Datum::from(n)
165            } else if n.is_nan() {
166                Datum::String("NaN")
167            } else if n.is_negative() {
168                Datum::String("-Infinity")
169            } else {
170                Datum::String("Infinity")
171            };
172            JsonbRef::from_datum(datum)
173        }
174        datum => JsonbRef::from_datum(datum),
175    }
176}
177
178#[sqlfunc]
179fn jsonb_array_length<'a>(a: JsonbRef<'a>) -> Result<Option<i32>, EvalError> {
180    match a.into_datum() {
181        Datum::List(list) => {
182            let count = list.iter().count();
183            match i32::try_from(count) {
184                Ok(len) => Ok(Some(len)),
185                Err(_) => Err(EvalError::Int32OutOfRange(count.to_string().into())),
186            }
187        }
188        _ => Ok(None),
189    }
190}
191
192#[sqlfunc]
193fn jsonb_typeof<'a>(a: JsonbRef<'a>) -> &'a str {
194    match a.into_datum() {
195        Datum::Map(_) => "object",
196        Datum::List(_) => "array",
197        Datum::String(_) => "string",
198        Datum::Numeric(_) => "number",
199        Datum::True | Datum::False => "boolean",
200        Datum::JsonNull => "null",
201        d => panic!("Not jsonb: {:?}", d),
202    }
203}
204
205#[sqlfunc]
206fn jsonb_strip_nulls<'a>(a: JsonbRef<'a>) -> Jsonb {
207    fn strip_nulls(a: Datum, row: &mut RowPacker) {
208        match a {
209            Datum::Map(dict) => row.push_dict_with(|row| {
210                for (k, v) in dict.iter() {
211                    match v {
212                        Datum::JsonNull => (),
213                        _ => {
214                            row.push(Datum::String(k));
215                            strip_nulls(v, row);
216                        }
217                    }
218                }
219            }),
220            Datum::List(list) => row.push_list_with(|row| {
221                for elem in list.iter() {
222                    strip_nulls(elem, row);
223                }
224            }),
225            _ => row.push(a),
226        }
227    }
228    let mut row = Row::default();
229    strip_nulls(a.into_datum(), &mut row.packer());
230    Jsonb::from_row(row)
231}
232
233#[sqlfunc]
234fn jsonb_pretty<'a>(a: JsonbRef<'a>) -> String {
235    let mut buf = String::new();
236    strconv::format_jsonb_pretty(&mut buf, a);
237    buf
238}
239
240/// Converts a JSONB `Datum` into a `u64`.
241fn jsonb_datum_to_u64<'a>(d: Datum<'a>) -> Result<u64, String> {
242    let Datum::Numeric(n) = d else {
243        return Err("expected numeric value".into());
244    };
245
246    let mut cx = numeric::cx_datum();
247    cx.try_into_u64(n.0)
248        .map_err(|_| format!("number out of u64 range: {n}"))
249}
250
251/// Decodes a JSONB object of shape `{"bitflags": <u64>}` into an `AclMode`.
252///
253/// Shared decoder for `parse_catalog_privileges` (which embeds the object as
254/// the `acl_mode` field of each privilege) and `parse_catalog_acl_mode` (which
255/// receives the object at the top level).
256fn jsonb_datum_to_acl_mode(d: Datum) -> Result<AclMode, String> {
257    let Datum::Map(dict) = d else {
258        return Err(format!("unexpected acl_mode: {d}"));
259    };
260    let mut bits = None;
261    for (key, val) in dict.iter() {
262        match key {
263            "bitflags" => bits = Some(jsonb_datum_to_u64(val)?),
264            other => return Err(format!("unexpected acl_mode field: {other}")),
265        }
266    }
267    let bits = bits.ok_or_else(|| "missing acl_mode bitflags".to_string())?;
268    AclMode::from_bits(bits).ok_or_else(|| format!("invalid acl_mode bitflags: {bits}"))
269}
270
271/// Converts a JSONB `Datum` into a `RoleId`.
272fn jsonb_datum_to_role_id(d: Datum) -> Result<RoleId, String> {
273    match d {
274        Datum::String("Public") => Ok(RoleId::Public),
275        Datum::String(other) => Err(format!("unexpected role ID variant: {other}")),
276        Datum::Map(dict) => {
277            let (key, val) = dict.iter().next().ok_or_else(|| "empty".to_string())?;
278            let n = jsonb_datum_to_u64(val)?;
279            match key {
280                "User" => Ok(RoleId::User(n)),
281                "System" => Ok(RoleId::System(n)),
282                "Predefined" => Ok(RoleId::Predefined(n)),
283                other => Err(format!("unexpected role ID variant: {other}")),
284            }
285        }
286        _ => Err("expected string or object".into()),
287    }
288}
289
290/// Converts a catalog JSON-serialized ID value into the appropriate string format.
291///
292/// Supports all of Materialize's various ID types of the form `<prefix><u64>`.
293#[sqlfunc]
294fn parse_catalog_id<'a>(a: JsonbRef<'a>) -> Result<String, EvalError> {
295    let parse = || match a.into_datum() {
296        // Unit variant, e.g. "Public"
297        Datum::String(variant) => match variant {
298            "Explain" => Ok("e".to_string()),
299            "Public" => Ok("p".to_string()),
300            other => Err(format!("unexpected ID variant: {other}")),
301        },
302        // Newtype variant, e.g. {"User": 1}
303        Datum::Map(dict) => {
304            let (key, val) = dict.iter().next().ok_or_else(|| "empty".to_string())?;
305            let prefix = match key {
306                "IntrospectionSourceIndex" => "si",
307                "Predefined" => "g",
308                "System" => "s",
309                "Transient" => "t",
310                "User" => "u",
311                other => return Err(format!("unexpected ID variant: {other}")),
312            };
313            let n = jsonb_datum_to_u64(val)?;
314            Ok(format!("{prefix}{n}"))
315        }
316        _ => Err("expected string or object".into()),
317    };
318
319    parse().map_err(|e| EvalError::InvalidCatalogJson(e.into()))
320}
321
322/// Converts a catalog JSON-serialized privilege array into an `mz_aclitem[]`.
323#[sqlfunc]
324fn parse_catalog_privileges<'a>(a: JsonbRef<'a>) -> Result<ArrayRustType<MzAclItem>, EvalError> {
325    let parse_one = |datum| match datum {
326        Datum::Map(dict) => {
327            let mut grantee = None;
328            let mut grantor = None;
329            let mut acl_mode = None;
330            for (key, val) in dict.iter() {
331                match key {
332                    "grantee" => {
333                        let id = jsonb_datum_to_role_id(val)?;
334                        grantee = Some(id);
335                    }
336                    "grantor" => {
337                        let id = jsonb_datum_to_role_id(val)?;
338                        grantor = Some(id);
339                    }
340                    "acl_mode" => {
341                        acl_mode = Some(jsonb_datum_to_acl_mode(val)?);
342                    }
343                    other => return Err(format!("unexpected privilege field: {other}")),
344                }
345            }
346            Ok(MzAclItem {
347                grantee: grantee.ok_or_else(|| format!("missing grantee: {dict:?}"))?,
348                grantor: grantor.ok_or_else(|| "missing grantor in privilege".to_string())?,
349                acl_mode: acl_mode.ok_or_else(|| "missing acl_mode in privilege".to_string())?,
350            })
351        }
352        other => Err(format!("expected object in array, found: {other}")),
353    };
354
355    let parse = || match a.into_datum() {
356        Datum::List(list) => {
357            let mut result = Vec::new();
358            for item in list.iter() {
359                result.push(parse_one(item)?);
360            }
361            Ok(result)
362        }
363        _ => Err("expected array".to_string()),
364    };
365
366    parse()
367        .map(ArrayRustType)
368        .map_err(|e| EvalError::InvalidCatalogJson(e.into()))
369}
370
371/// Converts a catalog JSON-serialized `AclMode` bitflags object into a
372/// PostgreSQL ACL char-code string (e.g. `{"bitflags": 514}` → `"ar"`).
373#[sqlfunc]
374fn parse_catalog_acl_mode<'a>(a: JsonbRef<'a>) -> Result<String, EvalError> {
375    jsonb_datum_to_acl_mode(a.into_datum())
376        .map(|mode| mode.to_string())
377        .map_err(|e| EvalError::InvalidCatalogJson(e.into()))
378}
379
380/// Parses a catalog `create_sql` string into a JSONB object.
381///
382/// The returned JSONB does not fully reflect the parsed SQL and instead contains only fields
383/// required by current callers.
384///
385// TODO: This function isn't parsing JSONB and therefore shouldn't live in the `jsonb` module.
386//       Consider moving all the `parse_catalog_*` functions into their own module.
387#[sqlfunc]
388fn parse_catalog_create_sql<'a>(a: &'a str) -> Result<Jsonb, EvalError> {
389    fn get_cluster_id(in_cluster: RawClusterName) -> Result<String, &'static str> {
390        match in_cluster {
391            RawClusterName::Resolved(s) => Ok(s),
392            RawClusterName::Unresolved(_) => Err("unresolved cluster name"),
393        }
394    }
395
396    fn get_item_id(item: RawItemName) -> Result<String, &'static str> {
397        match item {
398            RawItemName::Id(id, _, _) => Ok(id),
399            RawItemName::Name(_) => Err("unresolved item name"),
400        }
401    }
402
403    fn format_name<T: AstInfo>(fmt: &Format<T>) -> &'static str {
404        match fmt {
405            Format::Bytes => "bytes",
406            Format::Avro(_) => "avro",
407            Format::Protobuf(_) => "protobuf",
408            Format::Regex(_) => "regex",
409            Format::Csv { .. } => "csv",
410            Format::Json { .. } => "json",
411            Format::Text => "text",
412        }
413    }
414
415    let parse = || -> Result<serde_json::Value, String> {
416        let mut stmts = mz_sql_parser::parser::parse_statements(a)
417            .map_err(|e| format!("failed to parse create_sql: {e}"))?;
418        let stmt = match stmts.len() {
419            1 => stmts.remove(0).ast,
420            n => return Err(format!("expected a single statement, found {n}")),
421        };
422
423        let mut info = BTreeMap::<&str, serde_json::Value>::new();
424
425        use mz_sql_parser::ast::Statement::*;
426        let item_type = match stmt {
427            CreateSecret(_) => "secret",
428            CreateConnection(stmt) => {
429                let connection_type = stmt.connection_type.as_str();
430                info.insert("connection_type", json!(connection_type));
431
432                "connection"
433            }
434            CreateView(_) => "view",
435            CreateMaterializedView(stmt) => {
436                let Some(in_cluster) = stmt.in_cluster else {
437                    return Err("missing IN CLUSTER".into());
438                };
439                let cluster_id = match in_cluster {
440                    RawClusterName::Unresolved(ident) => ident.into_string(),
441                    RawClusterName::Resolved(s) => s,
442                };
443                info.insert("cluster_id", json!(cluster_id));
444
445                let mut definition = stmt.query.to_ast_string_stable();
446                definition.push(';');
447                info.insert("definition", json!(definition));
448
449                "materialized-view"
450            }
451            CreateTable(_) | CreateTableFromSource(_) => "table",
452            CreateSource(stmt) => {
453                let Some(in_cluster) = stmt.in_cluster else {
454                    return Err("missing IN CLUSTER".into());
455                };
456                let cluster_id = get_cluster_id(in_cluster)?;
457                info.insert("cluster_id", json!(cluster_id));
458
459                use mz_sql_parser::ast::CreateSourceConnection::*;
460                let (source_type, connection) = match stmt.connection {
461                    Kafka { connection, .. } => ("kafka", Some(connection)),
462                    Postgres { connection, .. } => ("postgres", Some(connection)),
463                    MySql { connection, .. } => ("mysql", Some(connection)),
464                    SqlServer { connection, .. } => ("sql-server", Some(connection)),
465                    LoadGenerator { .. } => ("load-generator", None),
466                };
467                info.insert("source_type", json!(source_type));
468                if let Some(conn) = connection {
469                    let conn_id = get_item_id(conn)?;
470                    info.insert("connection_id", json!(conn_id));
471                }
472
473                let is_debezium = matches!(
474                    stmt.envelope,
475                    Some(mz_sql_parser::ast::SourceEnvelope::Debezium)
476                );
477
478                if let Some(envelope) = stmt.envelope {
479                    use mz_sql_parser::ast::SourceEnvelope::*;
480                    let envelope_type = match envelope {
481                        None => "none",
482                        Debezium => "debezium",
483                        Upsert { .. } => "upsert",
484                        CdcV2 => "materialize",
485                    };
486                    info.insert("envelope_type", json!(envelope_type));
487                }
488
489                if let Some(format_spec) = stmt.format {
490                    match &format_spec {
491                        FormatSpecifier::Bare(fmt) => {
492                            // Debezium sources with a single format spec implicitly use
493                            // the same format for both key and value.
494                            if is_debezium {
495                                info.insert("key_format", json!(format_name(fmt)));
496                            }
497                            info.insert("value_format", json!(format_name(fmt)));
498                        }
499                        FormatSpecifier::KeyValue { key, value } => {
500                            info.insert("key_format", json!(format_name(key)));
501                            info.insert("value_format", json!(format_name(value)));
502                        }
503                    }
504                }
505
506                "source"
507            }
508            CreateWebhookSource(stmt) => {
509                if stmt.is_table {
510                    "table"
511                } else {
512                    info.insert("source_type", json!("webhook"));
513                    if let Some(in_cluster) = stmt.in_cluster {
514                        let cluster_id = get_cluster_id(in_cluster)?;
515                        info.insert("cluster_id", json!(cluster_id));
516                    }
517                    "source"
518                }
519            }
520            CreateSubsource(stmt) => {
521                use mz_sql_parser::ast::CreateSubsourceOptionName;
522                let is_progress = stmt
523                    .with_options
524                    .iter()
525                    .any(|o| matches!(o.name, CreateSubsourceOptionName::Progress));
526                let source_type = if is_progress { "progress" } else { "subsource" };
527                info.insert("source_type", json!(source_type));
528
529                if let Some(of_source) = stmt.of_source {
530                    let of_source_id = get_item_id(of_source)?;
531                    info.insert("of_source_id", json!(of_source_id));
532                }
533
534                "subsource"
535            }
536            CreateSink(_) => "sink",
537            CreateIndex(stmt) => {
538                let Some(in_cluster) = stmt.in_cluster else {
539                    return Err("missing IN CLUSTER".into());
540                };
541                let cluster_id = get_cluster_id(in_cluster)?;
542                info.insert("cluster_id", json!(cluster_id));
543                let on_id = get_item_id(stmt.on_name)?;
544                info.insert("on_id", json!(on_id));
545                "index"
546            }
547            CreateType(_) => "type",
548            _ => return Err("not a CREATE item statement".into()),
549        };
550        info.insert("type", json!(item_type));
551
552        let info = info.into_iter().map(|(k, v)| (k.to_string(), v)).collect();
553        Ok(info)
554    };
555
556    let val = parse().map_err(|e| EvalError::InvalidCatalogJson(e.into()))?;
557    let jsonb = Jsonb::from_serde_json(val).expect("valid JSONB");
558    Ok(jsonb)
559}
mz_expr/scalar/func/impls/jsonb.rs

mz_expr/scalar/func/impls/
jsonb.rs
