dre_core/
saved.rs

1//! Saved query results: each query's result as the exact Arrow IPC stream core received
2//! (`<name>.arrows`), plus an index (`results.json`). `dre run --output-json <DIR>` writes them
3//! beside the typed JSON; `dre run --from-results <DIR>` replays them to the format plugins
4//! without running any query. Types (decimals, timezones, nulls) stay exact, and the files open
5//! directly in DuckDB, Polars or pandas.
6
7use std::fs::File;
8use std::io::{BufReader, BufWriter};
9use std::path::Path;
10
11use arrow::datatypes::SchemaRef;
12use arrow::ipc::reader::{FileReader, StreamReader};
13use arrow::ipc::writer::{FileWriter, StreamWriter};
14use serde::{Deserialize, Serialize};
15
16/// The index file's name.
17pub const INDEX: &str = "results.json";
18pub const SCHEMA_VERSION: &str = "dre/saved-results/v1";
19
20/// `results.json`: the saved results of one Binding's run.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub struct Index {
23    pub schema_version: String,
24    /// The run that saved them.
25    pub run_id: String,
26    pub results: Vec<SavedResult>,
27}
28
29/// One query's saved result.
30#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
31pub struct SavedResult {
32    pub query: String,
33    /// Its tab name.
34    pub name: String,
35    pub connection: String,
36    pub rows: u64,
37    /// The SHA-256 of the query's rendered SQL, to warn when it changed since.
38    #[serde(default)]
39    pub sql_sha256: Option<String>,
40    /// The Arrow IPC stream, relative to the index.
41    pub file: String,
42    /// `secret_columns` saved as `*****`: a replay can't rebuild what they fed.
43    #[serde(default, skip_serializing_if = "Vec::is_empty")]
44    pub masked_columns: Vec<String>,
45}
46
47impl Index {
48    pub fn find(&self, query: &str) -> Option<&SavedResult> {
49        self.results.iter().find(|r| r.query == query)
50    }
51}
52
53/// Write the index into `dir`.
54pub fn write_index(dir: &Path, index: &Index) -> std::io::Result<()> {
55    crate::fs::write_atomic(&dir.join(INDEX), serde_json::to_string_pretty(index)?.as_bytes())
56}
57
58/// Read `dir`'s index. The error is a sentence.
59pub fn read_index(dir: &Path) -> Result<Index, String> {
60    let path = dir.join(INDEX);
61    let text = std::fs::read_to_string(&path).map_err(|e| format!("can't read {}: {e}", path.display()))?;
62    let index: Index =
63        serde_json::from_str(&text).map_err(|e| format!("{} is malformed: {e}", path.display()))?;
64    if index.schema_version != SCHEMA_VERSION {
65        return Err(format!(
66            "{} is `{}`, which this DRE doesn't read (it reads `{SCHEMA_VERSION}`)",
67            path.display(),
68            index.schema_version
69        ));
70    }
71    Ok(index)
72}
73
74/// Copy a spool (an Arrow IPC file) to `to` as an Arrow IPC stream.
75pub fn spool_to_stream(spool: &Path, to: &Path) -> Result<(), String> {
76    let reader = FileReader::try_new(File::open(spool).map_err(|e| e.to_string())?, None)
77        .map_err(|e| e.to_string())?;
78    let schema = reader.schema();
79    let mut w = StreamWriter::try_new(
80        BufWriter::new(File::create(to).map_err(|e| e.to_string())?),
81        &schema,
82    )
83    .map_err(|e| e.to_string())?;
84    for b in reader {
85        w.write(&b.map_err(|e| e.to_string())?)
86            .map_err(|e| e.to_string())?;
87    }
88    w.finish().map_err(|e| e.to_string())
89}
90
91/// Copy a saved stream back into a spool (an Arrow IPC file): its schema and row count.
92pub fn stream_to_spool(from: &Path, spool: &Path) -> Result<(SchemaRef, u64), String> {
93    let reader = StreamReader::try_new(
94        BufReader::new(File::open(from).map_err(|e| format!("can't read {}: {e}", from.display()))?),
95        None,
96    )
97    .map_err(|e| format!("{} isn't an Arrow IPC stream: {e}", from.display()))?;
98    let schema = reader.schema();
99    let mut w = FileWriter::try_new(File::create(spool).map_err(|e| e.to_string())?, &schema)
100        .map_err(|e| e.to_string())?;
101    let mut rows = 0u64;
102    for b in reader {
103        let b = b.map_err(|e| format!("{}: {e}", from.display()))?;
104        rows += b.num_rows() as u64;
105        w.write(&b).map_err(|e| e.to_string())?;
106    }
107    w.finish().map_err(|e| e.to_string())?;
108    Ok((schema, rows))
109}
110
111#[cfg(test)]
112mod tests {
113    use std::sync::Arc;
114
115    use arrow::array::{ArrayRef, Decimal128Array, RecordBatch, StringArray};
116
117    use super::*;
118
119    #[test]
120    fn a_result_round_trips_exactly() {
121        let dir = tempfile::tempdir().unwrap();
122        let batch = RecordBatch::try_from_iter([
123            (
124                "amount",
125                Arc::new(
126                    Decimal128Array::from(vec![Some(12345), None])
127                        .with_precision_and_scale(18, 2)
128                        .unwrap(),
129                ) as ArrayRef,
130            ),
131            (
132                "name",
133                Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef,
134            ),
135        ])
136        .unwrap();
137        let spool = dir.path().join("0.arrow");
138        let mut w = FileWriter::try_new(File::create(&spool).unwrap(), &batch.schema()).unwrap();
139        w.write(&batch).unwrap();
140        w.finish().unwrap();
141        let stream = dir.path().join("q.arrows");
142        spool_to_stream(&spool, &stream).unwrap();
143        let back = dir.path().join("1.arrow");
144        let (schema, rows) = stream_to_spool(&stream, &back).unwrap();
145        assert_eq!(rows, 2);
146        assert_eq!(schema, batch.schema());
147        let read: Vec<RecordBatch> = FileReader::try_new(File::open(&back).unwrap(), None)
148            .unwrap()
149            .map(Result::unwrap)
150            .collect();
151        assert_eq!(read, vec![batch]);
152    }
153}
dre_core/saved.rs

dre_core/
saved.rs
