grafbase_sdk/
jq_selection.rs1use 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
31pub struct JqSelection {
38 arena: Arena,
39 inputs: RcIter<Empty<Result<Val, String>>>,
41 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 pub fn new() -> Self {
70 Self::default()
71 }
72
73 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}