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

grafbase_sdk/
jq_selection.rs
