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}
43
44impl Index {
45    pub fn find(&self, query: &str) -> Option<&SavedResult> {
46        self.results.iter().find(|r| r.query == query)
47    }
48}
49
50/// Write the index into `dir`.
51pub fn write_index(dir: &Path, index: &Index) -> std::io::Result<()> {
52    crate::fs::write_atomic(&dir.join(INDEX), serde_json::to_string_pretty(index)?.as_bytes())
53}
54
55/// Read `dir`'s index. The error is a sentence.
56pub fn read_index(dir: &Path) -> Result<Index, String> {
57    let path = dir.join(INDEX);
58    let text = std::fs::read_to_string(&path).map_err(|e| format!("can't read {}: {e}", path.display()))?;
59    let index: Index =
60        serde_json::from_str(&text).map_err(|e| format!("{} is malformed: {e}", path.display()))?;
61    if index.schema_version != SCHEMA_VERSION {
62        return Err(format!(
63            "{} is `{}`, which this DRE doesn't read (it reads `{SCHEMA_VERSION}`)",
64            path.display(),
65            index.schema_version
66        ));
67    }
68    Ok(index)
69}
70
71/// Copy a spool (an Arrow IPC file) to `to` as an Arrow IPC stream.
72pub fn spool_to_stream(spool: &Path, to: &Path) -> Result<(), String> {
73    let reader = FileReader::try_new(File::open(spool).map_err(|e| e.to_string())?, None)
74        .map_err(|e| e.to_string())?;
75    let schema = reader.schema();
76    let mut w = StreamWriter::try_new(
77        BufWriter::new(File::create(to).map_err(|e| e.to_string())?),
78        &schema,
79    )
80    .map_err(|e| e.to_string())?;
81    for b in reader {
82        w.write(&b.map_err(|e| e.to_string())?)
83            .map_err(|e| e.to_string())?;
84    }
85    w.finish().map_err(|e| e.to_string())
86}
87
88/// Copy a saved stream back into a spool (an Arrow IPC file): its schema and row count.
89pub fn stream_to_spool(from: &Path, spool: &Path) -> Result<(SchemaRef, u64), String> {
90    let reader = StreamReader::try_new(
91        BufReader::new(File::open(from).map_err(|e| format!("can't read {}: {e}", from.display()))?),
92        None,
93    )
94    .map_err(|e| format!("{} isn't an Arrow IPC stream: {e}", from.display()))?;
95    let schema = reader.schema();
96    let mut w = FileWriter::try_new(File::create(spool).map_err(|e| e.to_string())?, &schema)
97        .map_err(|e| e.to_string())?;
98    let mut rows = 0u64;
99    for b in reader {
100        let b = b.map_err(|e| format!("{}: {e}", from.display()))?;
101        rows += b.num_rows() as u64;
102        w.write(&b).map_err(|e| e.to_string())?;
103    }
104    w.finish().map_err(|e| e.to_string())?;
105    Ok((schema, rows))
106}
107
108#[cfg(test)]
109mod tests {
110    use std::sync::Arc;
111
112    use arrow::array::{ArrayRef, Decimal128Array, RecordBatch, StringArray};
113
114    use super::*;
115
116    #[test]
117    fn a_result_round_trips_exactly() {
118        let dir = tempfile::tempdir().unwrap();
119        let batch = RecordBatch::try_from_iter([
120            (
121                "amount",
122                Arc::new(
123                    Decimal128Array::from(vec![Some(12345), None])
124                        .with_precision_and_scale(18, 2)
125                        .unwrap(),
126                ) as ArrayRef,
127            ),
128            (
129                "name",
130                Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef,
131            ),
132        ])
133        .unwrap();
134        let spool = dir.path().join("0.arrow");
135        let mut w = FileWriter::try_new(File::create(&spool).unwrap(), &batch.schema()).unwrap();
136        w.write(&batch).unwrap();
137        w.finish().unwrap();
138        let stream = dir.path().join("q.arrows");
139        spool_to_stream(&spool, &stream).unwrap();
140        let back = dir.path().join("1.arrow");
141        let (schema, rows) = stream_to_spool(&stream, &back).unwrap();
142        assert_eq!(rows, 2);
143        assert_eq!(schema, batch.schema());
144        let read: Vec<RecordBatch> = FileReader::try_new(File::open(&back).unwrap(), None)
145            .unwrap()
146            .map(Result::unwrap)
147            .collect();
148        assert_eq!(read, vec![batch]);
149    }
150}
dre_core/saved.rs

dre_core/
saved.rs
