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::{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<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 hasher = |val: &_| hasher.hash_one(val);
85        let hash = hasher(selection);
86
87        if !self.selection_cache.find(hash, |_| true).is_some() {
88            self.create_filter(selection, hash)?;
89        }
90
91        let filter = self.find_filter(hash).unwrap();
92        let filtered = filter.run((Ctx::new([], &self.inputs), Val::from(data)));
93
94        Ok(filtered.map(|v| match v {
95            Ok(val) => Ok(serde_json::Value::from(val)),
96            Err(e) => Err(anyhow::anyhow!("{e}")),
97        }))
98    }
99
100    fn find_filter(&self, hash: u64) -> Option<&Filter<Native<Val>>> {
101        self.selection_cache
102            .find(hash, |_| true)
103            .and_then(|i| self.filters.get(*i))
104    }
105
106    fn create_filter(&mut self, selection: &str, hash: u64) -> anyhow::Result<()> {
107        let program = File {
108            code: selection,
109            path: (),
110        };
111
112        let loader = Loader::new(jaq_std::defs().chain(jaq_json::defs()));
113
114        let modules = loader.load(&self.arena, program).map_err(|e| {
115            let error = e.first().map(|e| e.0.code).unwrap_or_default();
116            anyhow::anyhow!("The selection is not valid jq syntax: `{error}`")
117        })?;
118
119        let filter = Compiler::default()
120            .with_funs(jaq_std::funs().chain(jaq_json::funs()))
121            .compile(modules)
122            .map_err(|e| {
123                let error = e.first().map(|e| e.0.code).unwrap_or_default();
124                anyhow::anyhow!("The selection is not valid jq syntax: `{error}`")
125            })?;
126
127        self.filters.push(filter);
128
129        let index = self.filters.len() - 1;
130        let hasher = DefaultHashBuilder::default();
131        let hasher = |val: &_| hasher.hash_one(val);
132
133        self.selection_cache.insert_unique(hash, index, hasher);
134
135        Ok(())
136    }
137}
grafbase_sdk/jq_selection.rs

grafbase_sdk/
jq_selection.rs
