mz_expr/scalar/func/
jq.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
10//! The `jq` function, which applies a jq program to a `jsonb` value.
11//!
12//! Programs run on jaq, a Rust implementation of jq. The available filters
13//! are jaq's standard library minus those whose result depends on the
14//! machine or the moment (`now`, `env`, `localtime`, `input`, ...), so that
15//! every replica computes the same answer. Module imports are not supported.
16//!
17//! NOTE: jaq has no evaluation budget. A program that loops forever stalls
18//! the worker evaluating it, and a program that recurses without bound
19//! overflows that worker's stack.
20
21use std::cmp::Ordering;
22use std::fmt;
23use std::hash::{Hash, Hasher};
24use std::str::FromStr;
25use std::sync::Arc;
26
27use jaq_core::compile::Undefined;
28use jaq_core::load::{self, Arena, File, Loader};
29use jaq_core::native::{Fun, run, v};
30use jaq_core::{Compiler, Ctx, Exn, Vars, data};
31use jaq_json::{Num, Val};
32use mz_expr_derive::sqlfunc;
33use mz_ore::str::StrExt;
34use mz_repr::Datum;
35use mz_repr::adt::jsonb::{Jsonb, JsonbRef};
36use serde::{Deserialize, Serialize};
37
38use crate::EvalError;
39
40type Data = data::JustLut<Val>;
41type Filter = jaq_core::Filter<Data>;
42
43/// Native filters withheld from programs, because their results depend on the
44/// environment, they read input other than the argument, or they write to the
45/// process's log or stderr.
46const WITHHELD: &[&str] = &[
47    "env",
48    "now",
49    "localtime",
50    "strflocaltime",
51    "input",
52    "inputs",
53    // Resolves time zone names against the host's time zone database.
54    "strptime",
55    "debug_empty",
56    "stderr_empty",
57];
58
59/// A compiled jq program.
60///
61/// Compares, orders, hashes, and serializes as its source text, and compiles
62/// again on deserialization.
63#[derive(Clone, Serialize, Deserialize)]
64#[serde(into = "String", try_from = "String")]
65pub struct JqProgram {
66    source: String,
67    filter: Arc<Filter>,
68}
69
70impl JqProgram {
71    /// Compiles `source`, or describes why it does not compile.
72    pub fn new(source: &str) -> Result<Self, String> {
73        Ok(JqProgram {
74            source: source.to_owned(),
75            filter: Arc::new(compile(source)?),
76        })
77    }
78
79    /// The program's source text.
80    pub fn source(&self) -> &str {
81        &self.source
82    }
83
84    /// Applies the program to `input`.
85    ///
86    /// A program that produces no value yields `None`. A program that
87    /// produces more than one value is an error, as jq's `[...]` collects a
88    /// stream into a single array.
89    pub fn eval(&self, input: JsonbRef<'_>) -> Result<Option<Jsonb>, EvalError> {
90        let input = to_val(input.into_datum());
91        let ctx = Ctx::<Data>::new(&self.filter.lut, Vars::new([]));
92        let mut outputs = self.filter.id.run((ctx, input));
93        let output = match outputs.next() {
94            None => return Ok(None),
95            Some(output) => output.map_err(exception)?,
96        };
97        if let Some(next) = outputs.next() {
98            next.map_err(exception)?;
99            return Err(error("program produced more than one value"));
100        }
101        // jq values are a superset of JSON, so output goes through the
102        // canonical `jsonb` text parser, which rejects anything that is not.
103        let output = Jsonb::from_str(&output.to_string())
104            .map_err(|e| error(format!("output is not JSON: {e}")))?;
105        Ok(Some(output))
106    }
107}
108
109impl From<JqProgram> for String {
110    fn from(program: JqProgram) -> String {
111        program.source
112    }
113}
114
115impl TryFrom<String> for JqProgram {
116    type Error = String;
117    fn try_from(source: String) -> Result<Self, String> {
118        JqProgram::new(&source)
119    }
120}
121
122impl PartialEq for JqProgram {
123    fn eq(&self, other: &Self) -> bool {
124        self.source == other.source
125    }
126}
127
128impl Eq for JqProgram {}
129
130impl PartialOrd for JqProgram {
131    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
132        Some(self.cmp(other))
133    }
134}
135
136impl Ord for JqProgram {
137    fn cmp(&self, other: &Self) -> Ordering {
138        self.source.cmp(&other.source)
139    }
140}
141
142impl Hash for JqProgram {
143    fn hash<H: Hasher>(&self, state: &mut H) {
144        self.source.hash(state);
145    }
146}
147
148impl fmt::Debug for JqProgram {
149    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
150        f.debug_tuple("JqProgram").field(&self.source).finish()
151    }
152}
153
154/// Applies a literal jq program to a `jsonb` value. See [`JqProgram::eval`].
155#[derive(
156    Ord,
157    PartialOrd,
158    Clone,
159    Debug,
160    Eq,
161    PartialEq,
162    Serialize,
163    Deserialize,
164    Hash
165)]
166pub struct Jq(pub JqProgram);
167
168#[sqlfunc(Jq, skip_display = true)]
169fn jq<'a>(&self, input: JsonbRef<'a>) -> Result<Option<Jsonb>, EvalError> {
170    self.0.eval(input)
171}
172
173impl fmt::Display for Jq {
174    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
175        write!(f, "jq[{}]", self.0.source().escaped())
176    }
177}
178
179fn compile(program: &str) -> Result<Filter, String> {
180    let defs = jaq_core::defs()
181        .chain(jaq_std::defs())
182        .chain(jaq_json::defs());
183    // The standard library defines `debug` and `stderr` in terms of these,
184    // so they stay defined, as no-ops.
185    let silenced: [Fun<Data>; 2] = [
186        run(("debug_empty", v(0), |_| Box::new(std::iter::empty()))),
187        run(("stderr_empty", v(0), |_| Box::new(std::iter::empty()))),
188    ];
189    let funs = jaq_core::funs()
190        .chain(jaq_std::funs())
191        .chain(jaq_json::funs())
192        .filter(|(name, _, _)| !WITHHELD.contains(name))
193        .chain(silenced);
194
195    let arena = Arena::default();
196    let modules = Loader::new(defs)
197        .load(
198            &arena,
199            File {
200                code: program,
201                path: (),
202            },
203        )
204        .map_err(|errs| {
205            let msgs = errs.into_iter().flat_map(|(_file, err)| match err {
206                load::Error::Io(errs) => errs
207                    .into_iter()
208                    .map(|(path, msg)| format!("cannot import {path}: {msg}"))
209                    .collect(),
210                load::Error::Lex(errs) => errs
211                    .into_iter()
212                    .map(|(expect, at)| syntax_error(program, expect.as_str(), at))
213                    .collect(),
214                load::Error::Parse(errs) => errs
215                    .into_iter()
216                    .map(|(expect, at)| syntax_error(program, expect.as_str(), at))
217                    .collect::<Vec<_>>(),
218            });
219            msgs.collect::<Vec<_>>().join("; ")
220        })?;
221    Compiler::default()
222        .with_funs(funs)
223        .compile(modules)
224        .map_err(|errs| {
225            let msgs = errs.into_iter().flat_map(|(_file, errs)| errs);
226            let msgs = msgs.map(|(name, undefined)| match undefined {
227                Undefined::Filter(arity) => format!("undefined filter {name}/{arity}"),
228                other => format!("undefined {} {name}", other.as_str()),
229            });
230            msgs.collect::<Vec<_>>().join("; ")
231        })
232}
233
234/// Describes a syntax error at `at`, a suffix of `program`.
235fn syntax_error(program: &str, expected: &str, at: &str) -> String {
236    if at.is_empty() {
237        format!("expected {expected} at end of program")
238    } else {
239        let offset = load::span(program, at).start;
240        format!("expected {expected} at offset {offset}")
241    }
242}
243
244fn exception(exn: Exn<Val>) -> EvalError {
245    match exn.get_err() {
246        Ok(err) => error(err),
247        Err(exn) => match exn.get_halt() {
248            Ok(code) => error(format!("program halted with exit code {code}")),
249            Err(_) => error("program raised an unexpected exception"),
250        },
251    }
252}
253
254fn error(msg: impl std::fmt::Display) -> EvalError {
255    EvalError::InvalidParameterValue(format!("jq: {msg}").into())
256}
257
258/// Converts a `jsonb` datum into a jaq value.
259///
260/// Numbers convert as jaq's own JSON reader converts their text: integers to
261/// `Num::Int` or `Num::BigInt`, and everything else to `Num::Dec`, which keeps
262/// the decimal text until arithmetic turns it into a float.
263fn to_val(datum: Datum<'_>) -> Val {
264    mz_ore::stack::maybe_grow(|| match datum {
265        Datum::JsonNull => Val::Null,
266        Datum::True => Val::Bool(true),
267        Datum::False => Val::Bool(false),
268        Datum::Numeric(n) => {
269            // Standard notation has no exponent, so an integer has no `.`.
270            let text = n.into_inner().to_standard_notation_string();
271            let num = if text.contains('.') {
272                None
273            } else {
274                Num::from_str_radix(&text, 10)
275            };
276            Val::Num(num.unwrap_or_else(|| Num::Dec(jaq_json::Rc::new(text))))
277        }
278        Datum::String(s) => Val::utf8_str(s.as_bytes().to_vec()),
279        Datum::List(list) => Val::Arr(jaq_json::Rc::new(list.iter().map(to_val).collect())),
280        Datum::Map(dict) => Val::obj(
281            dict.iter()
282                .map(|(k, v)| (Val::utf8_str(k.as_bytes().to_vec()), to_val(v)))
283                .collect(),
284        ),
285        d => unreachable!("not a jsonb datum: {d:?}"),
286    })
287}
288
289#[cfg(test)]
290mod tests;
mz_expr/scalar/func/jq.rs

mz_expr/scalar/func/
jq.rs
