dre_core/
for_each.rs

1//! Per-row outputs (`for_each`): the pieces that don't need a run. The report's queries run once;
2//! the driver query gives one row per copy of the output, and every other query of the output is
3//! split by the same key, so each copy gets only its own slice. See `run.rs` for the run itself.
4
5use std::collections::{BTreeMap, HashMap};
6use std::sync::Arc;
7
8use arrow::array::{RecordBatch, UInt32Array};
9use arrow::datatypes::SchemaRef;
10use arrow::util::display::{ArrayFormatter, FormatOptions};
11
12/// Joins the parts of a composite key in an index (never in a value).
13const SEP: char = '\u{1f}';
14
15/// A query's rows, in one batch, with its key column positions.
16pub struct Table {
17    pub name: String,
18    pub batch: RecordBatch,
19}
20
21impl Table {
22    /// One batch from a result's batches (an empty one keeps the schema).
23    pub fn new(name: &str, schema: SchemaRef, batches: &[RecordBatch]) -> Result<Table, String> {
24        let batch = arrow::compute::concat_batches(&schema, batches).map_err(|e| e.to_string())?;
25        Ok(Table {
26            name: name.to_string(),
27            batch,
28        })
29    }
30
31    fn key_columns(&self, key: &[String]) -> Result<Vec<usize>, String> {
32        key.iter()
33            .map(|k| {
34                self.batch.schema().index_of(k).map_err(|_| {
35                    format!(
36                        "`{}` has no column `{k}`, the `for_each` key; every query of a per-row output needs the key columns (its columns: {})",
37                        self.name,
38                        self.batch
39                            .schema()
40                            .fields()
41                            .iter()
42                            .map(|f| f.name().as_str())
43                            .collect::<Vec<_>>()
44                            .join(", ")
45                    )
46                })
47            })
48            .collect()
49    }
50
51    /// Every row's key: (index form, shown form `a_b`).
52    fn keys(&self, cols: &[usize]) -> Result<Vec<(String, String)>, String> {
53        let opts = FormatOptions::default().with_null("");
54        let fmts: Vec<ArrayFormatter<'_>> = cols
55            .iter()
56            .map(|&c| ArrayFormatter::try_new(self.batch.column(c).as_ref(), &opts))
57            .collect::<Result<_, _>>()
58            .map_err(|e| e.to_string())?;
59        Ok((0..self.batch.num_rows())
60            .map(|r| {
61                let parts: Vec<String> = fmts.iter().map(|f| f.value(r).to_string()).collect();
62                (parts.join(&SEP.to_string()), parts.join("_"))
63            })
64            .collect())
65    }
66}
67
68/// The population of a per-row output.
69pub struct Population {
70    pub driver: Table,
71    /// Each driver row's key, shown (`a_b` for a composite key), in driver order.
72    pub keys: Vec<String>,
73    /// Each data query, and its rows by key (index form).
74    pub data: Vec<(Table, HashMap<String, Vec<u32>>)>,
75    /// The driver rows' keys (index form), in driver order.
76    index_keys: Vec<String>,
77    /// Rows of data queries whose key isn't in the driver: how many, how many keys, and up to
78    /// 10 of them.
79    pub orphan_rows: u64,
80    pub orphan_keys: u64,
81    pub orphan_samples: Vec<String>,
82}
83
84impl Population {
85    /// Index the driver and each data query by `key`. A duplicate driver key is an error.
86    pub fn build(driver: Table, data: Vec<Table>, key: &[String]) -> Result<Population, String> {
87        let cols = driver.key_columns(key)?;
88        let keys = driver.keys(&cols)?;
89        let mut seen: HashMap<&str, usize> = HashMap::new();
90        for (i, (k, shown)) in keys.iter().enumerate() {
91            if let Some(first) = seen.insert(k.as_str(), i) {
92                return Err(format!(
93                    "the driver `{}` has the key `{shown}` twice (rows {} and {}); a per-row output needs one driver row per key",
94                    driver.name,
95                    first + 1,
96                    i + 1
97                ));
98            }
99        }
100        let mut orphan_rows = 0;
101        let mut orphans: std::collections::HashSet<String> = std::collections::HashSet::new();
102        let mut orphan_samples = Vec::new();
103        let mut indexed = Vec::new();
104        for t in data {
105            let cols = t.key_columns(key)?;
106            let mut index: HashMap<String, Vec<u32>> = HashMap::new();
107            for (r, (k, shown)) in t.keys(&cols)?.into_iter().enumerate() {
108                if !seen.contains_key(k.as_str()) {
109                    orphan_rows += 1;
110                    orphans.insert(k.clone());
111                    if orphan_samples.len() < 10 && !orphan_samples.contains(&shown) {
112                        orphan_samples.push(shown);
113                    }
114                    continue;
115                }
116                index.entry(k).or_default().push(r as u32);
117            }
118            indexed.push((t, index));
119        }
120        let (index_keys, keys) = keys.into_iter().unzip();
121        Ok(Population {
122            driver,
123            keys,
124            data: indexed,
125            index_keys,
126            orphan_rows,
127            orphan_keys: orphans.len() as u64,
128            orphan_samples,
129        })
130    }
131
132    pub fn len(&self) -> usize {
133        self.keys.len()
134    }
135
136    pub fn is_empty(&self) -> bool {
137        self.keys.is_empty()
138    }
139
140    /// Row `i`'s slice of each data query, in order.
141    pub fn slices(&self, i: usize) -> Result<Vec<RecordBatch>, String> {
142        let k = &self.index_keys[i];
143        self.data
144            .iter()
145            .map(|(t, index)| {
146                let rows = index.get(k).map(Vec::as_slice).unwrap_or(&[]);
147                take(&t.batch, rows)
148            })
149            .collect()
150    }
151
152    /// How many rows of data query `q` row `i` has.
153    pub fn slice_len(&self, i: usize, q: usize) -> u64 {
154        self.data[q]
155            .1
156            .get(&self.index_keys[i])
157            .map_or(0, |r| r.len() as u64)
158    }
159
160    /// Row `i` of the driver, as a one-row batch.
161    pub fn driver_row(&self, i: usize) -> RecordBatch {
162        self.driver.batch.slice(i, 1)
163    }
164}
165
166fn take(batch: &RecordBatch, rows: &[u32]) -> Result<RecordBatch, String> {
167    let idx = UInt32Array::from(rows.to_vec());
168    arrow::compute::take_record_batch(batch, &idx).map_err(|e| e.to_string())
169}
170
171/// A key as a file name: letters, digits, `.`, `-` and `_` kept, anything else `_`; never empty,
172/// `.` or `..`, at most 100 characters.
173pub fn safe_name(key: &str) -> String {
174    let mut s: String = key
175        .chars()
176        .map(|c| {
177            if c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | '_') {
178                c
179            } else {
180                '_'
181            }
182        })
183        .take(100)
184        .collect();
185    if s.is_empty() || s.chars().all(|c| c == '.') {
186        s = format!("_{s}");
187    }
188    s
189}
190
191/// Files of two rows that would collide: (file, first key, second key).
192pub fn collision<'a>(
193    names: impl IntoIterator<Item = (&'a str, &'a str)>,
194) -> Option<(String, String, String)> {
195    let mut seen: BTreeMap<String, &str> = BTreeMap::new();
196    for (key, file) in names {
197        if let Some(prev) = seen.insert(file.to_lowercase(), key) {
198            return Some((file.to_string(), prev.to_string(), key.to_string()));
199        }
200    }
201    None
202}
203
204/// A per-row value as logs and results show it: an email keeps its first character and domain
205/// (`j***@example.com`), anything else is `*****`.
206pub fn mask_value(v: &str) -> String {
207    if let Some((user, domain)) = v.split_once('@')
208        && !user.is_empty()
209        && domain.contains('.')
210        && !v.contains(char::is_whitespace)
211    {
212        let first: String = user.chars().take(1).collect();
213        return format!("{first}***@{domain}");
214    }
215    crate::secrets::MASK.to_string()
216}
217
218/// Masks a row's values wherever they appear in a text (a rendered path, a plugin's error).
219#[derive(Debug, Clone, Default)]
220pub struct RowMask {
221    /// (value, its mask), longest first.
222    values: Arc<Vec<(String, String)>>,
223}
224
225impl RowMask {
226    /// The row's string values worth hiding: every `each.*` value but the key's, at least three
227    /// characters long (shorter ones would mask parts of unrelated words).
228    pub fn new(each: &serde_json::Map<String, serde_json::Value>, key: &[String]) -> RowMask {
229        let mut values: Vec<(String, String)> = each
230            .iter()
231            .filter(|(k, _)| !key.contains(k))
232            .filter_map(|(_, v)| match v {
233                serde_json::Value::String(s) => Some(s.clone()),
234                serde_json::Value::Number(n) => Some(n.to_string()),
235                _ => None,
236            })
237            .filter(|s| s.chars().count() >= 3)
238            .map(|s| {
239                let m = mask_value(&s);
240                (s, m)
241            })
242            .collect();
243        values.sort_by_key(|(s, _)| std::cmp::Reverse(s.len()));
244        values.dedup();
245        RowMask {
246            values: Arc::new(values),
247        }
248    }
249
250    pub fn apply(&self, text: &str) -> String {
251        let mut out = text.to_string();
252        for (v, m) in self.values.iter() {
253            if out.contains(v.as_str()) {
254                out = out.replace(v.as_str(), m);
255            }
256        }
257        out
258    }
259
260    /// Mask every string in a JSON value.
261    pub fn apply_json(&self, v: &serde_json::Value) -> serde_json::Value {
262        match v {
263            serde_json::Value::String(s) => serde_json::Value::String(self.apply(s)),
264            serde_json::Value::Array(a) => {
265                serde_json::Value::Array(a.iter().map(|x| self.apply_json(x)).collect())
266            }
267            serde_json::Value::Object(o) => {
268                serde_json::Value::Object(o.iter().map(|(k, x)| (k.clone(), self.apply_json(x))).collect())
269            }
270            other => other.clone(),
271        }
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278    use arrow::array::{ArrayRef, Int64Array, StringArray};
279
280    fn table(name: &str, ids: &[i64], other: &[&str]) -> Table {
281        let b = RecordBatch::try_from_iter([
282            ("member_id", Arc::new(Int64Array::from(ids.to_vec())) as ArrayRef),
283            ("v", Arc::new(StringArray::from(other.to_vec())) as ArrayRef),
284        ])
285        .unwrap();
286        Table::new(name, b.schema(), &[b]).unwrap()
287    }
288
289    #[test]
290    fn slices_by_key_and_counts_orphans() {
291        let driver = table("members", &[1, 2, 3], &["a", "b", "c"]);
292        let tx = table("tx", &[2, 1, 2, 9, 9], &["x", "y", "z", "o1", "o2"]);
293        let p = Population::build(driver, vec![tx], &["member_id".into()]).unwrap();
294        assert_eq!(p.keys, vec!["1", "2", "3"]);
295        let s = p.slices(1).unwrap();
296        let v = s[0].column(1).as_any().downcast_ref::<StringArray>().unwrap();
297        assert_eq!(v.iter().map(Option::unwrap).collect::<Vec<_>>(), vec!["x", "z"]);
298        assert_eq!(p.slices(2).unwrap()[0].num_rows(), 0);
299        assert_eq!(
300            (p.orphan_rows, p.orphan_samples.clone()),
301            (2, vec!["9".to_string()])
302        );
303    }
304
305    #[test]
306    fn composite_and_duplicate_keys() {
307        let d = table("members", &[1, 1], &["a", "b"]);
308        let p = Population::build(d, vec![], &["member_id".into(), "v".into()]).unwrap();
309        assert_eq!(p.keys, vec!["1_a", "1_b"]);
310        let d = table("members", &[1, 1], &["a", "a"]);
311        let e = Population::build(d, vec![], &["member_id".into()]).err().unwrap();
312        assert!(e.contains("has the key `1` twice (rows 1 and 2)"), "{e}");
313        let d = table("members", &[1], &["a"]);
314        let e = Population::build(d, vec![], &["nope".into()]).err().unwrap();
315        assert!(e.contains("has no column `nope`"), "{e}");
316    }
317
318    #[test]
319    fn file_names_are_safe_and_collisions_found() {
320        assert_eq!(safe_name("ab/../c d"), "ab_.._c_d");
321        assert_eq!(safe_name(".."), "_..");
322        assert_eq!(safe_name(""), "_");
323        assert_eq!(
324            collision([("a b", "a_b.csv"), ("a/b", "a_b.csv")]),
325            Some(("a_b.csv".into(), "a b".into(), "a/b".into()))
326        );
327        assert_eq!(collision([("a", "a.csv"), ("b", "b.csv")]), None);
328    }
329
330    #[test]
331    fn masking() {
332        assert_eq!(mask_value("jane.doe@example.com"), "j***@example.com");
333        assert_eq!(mask_value("Jane Doe"), "*****");
334        let each = serde_json::json!({"member_id": 12345, "email": "jane@example.com", "name": "Jane Doe", "x": "ab"});
335        let m = RowMask::new(each.as_object().unwrap(), &["member_id".into()]);
336        assert_eq!(
337            m.apply("to jane@example.com for Jane Doe, member 12345, ab"),
338            "to j***@example.com for *****, member 12345, ab"
339        );
340    }
341}
dre_core/for_each.rs

dre_core/
for_each.rs
