grafbase_sdk/
jq_selection.rs

1//! A module for performing JQ filter selections on JSON data.
2//!
3//! Enable feature `jq-selection` to use this module.
4//!
5//! This module provides functionality to parse, compile and execute JQ filter
6//! expressions against JSON data. It internally caches compiled filters to avoid
7//! recompilation overhead.
8//!
9//! # Examples
10//!
11//! ```rust
12//! # use grafbase_sed::jq_selection::JqSelection;
13//!
14//! let mut jq = JqSelection::new();
15//! let data = serde_json::json!({"name": "Jane", "age": 25});
16//! let results = jq.select(".name", data).unwrap();
17//!
18//! assert_eq!(results, serde_json::json!("Jane"));
19//! ```
20
21use std::iter::Empty;
22
23use core::hash::BuildHasher;
24use hashbrown::{hash_table::Entry, DefaultHashBuilder, HashTable};
25use jaq_core::{
26    load::{Arena, File, Loader},
27    Compiler, Ctx, Filter, Native, RcIter,
28};
29use jaq_json::Val;
30
31/// A struct that holds JQ filter selections
32///
33/// Use it to select data from a JSON object using JQ syntax. Caches the previously compiled filters,
34/// and reuses them to avoid recompiling the same filter multiple times.
35///
36/// You are supposed to store this struct in your extension and reuse it across multiple requests.
37pub struct JqSelection {
38    arena: Arena,
39    // (╯° · °)╯︵ ┻━┻
40    inputs: RcIter<Empty<Result<Val, String>>>,
41    // ┬┴┬┴┤
42    // ┬┴┬┴┤ ͡°)
43    // ┬┴┬┴┤ ͜ʖ ͡°)
44    // ┬┴┬┴┤ ͡° ͜ʖ ͡°)
45    // ┬┴┬┴┤ ͡° ͜ʖ ͡~)
46    // ┬┴┬┴┤ ͡° ͜ʖ ͡°)
47    // ┬┴┬┴┤ ͜ʖ ͡°)
48    // ┬┴┬┴┤ ͡°)
49    // ┬┴┬┴┤
50    selection_cache: HashTable<(String, usize)>,
51    filters: Vec<Filter<Native<Val>>>,
52}
53
54impl Default for JqSelection {
55    fn default() -> Self {
56        Self {
57            arena: Arena::default(),
58            inputs: RcIter::new(core::iter::empty()),
59            selection_cache: HashTable::new(),
60            filters: Vec::new(),
61        }
62    }
63}
64
65impl JqSelection {
66    /// Creates a new instance of [`JqSelection`].
67    ///
68    /// Creates an empty cache of compiled filters.
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    /// Selects data from a JSON value using a JQ filter.
74    ///
75    /// This method takes a JQ selection filter string and a JSON value, applies the
76    /// filter, and returns an iterator of the results. The filter is compiled and cached
77    /// for reuse on subsequent calls with the same filter string.
78    pub fn select(
79        &mut self,
80        selection: &str,
81        data: serde_json::Value,
82    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<serde_json::Value>> + '_> {
83        let hasher = DefaultHashBuilder::default();
84        let hash = hasher.hash_one(selection);
85        let hasher = |val: &(String, usize)| hasher.hash_one(&val.0);
86
87        let idx = match self
88            .selection_cache
89            .entry(hash, |(key, _)| key.as_str() == selection, hasher)
90        {
91            Entry::Occupied(entry) => entry.get().1,
92            Entry::Vacant(vacant_entry) => {
93                let program = File {
94                    code: selection,
95                    path: (),
96                };
97
98                let loader = Loader::new(jaq_std::defs().chain(jaq_json::defs()));
99
100                let modules = loader.load(&self.arena, program).map_err(|e| {
101                    let error = e.first().map(|e| e.0.code).unwrap_or_default();
102                    anyhow::anyhow!("The selection is not valid jq syntax: `{error}`")
103                })?;
104
105                let filter = Compiler::default()
106                    .with_funs(jaq_std::funs().chain(jaq_json::funs()))
107                    .compile(modules)
108                    .map_err(|e| {
109                        let error = e.first().map(|e| e.0.code).unwrap_or_default();
110                        anyhow::anyhow!("The selection is not valid jq syntax: `{error}`")
111                    })?;
112
113                self.filters.push(filter);
114
115                let index = self.filters.len() - 1;
116                vacant_entry.insert((selection.to_string(), index));
117
118                index
119            }
120        };
121
122        let filter = &self.filters[idx];
123        let filtered = filter.run((Ctx::new([], &self.inputs), Val::from(data)));
124
125        Ok(filtered.map(|v| match v {
126            Ok(val) => Ok(serde_json::Value::from(val)),
127            Err(e) => Err(anyhow::anyhow!("{e}")),
128        }))
129    }
130}