Skip to main content

dre_core/
message.rs

1//! The built-in `message` format: a Binding's query results rendered through Jinja into a short
2//! headline, a title plus text in portable Markdown ([`dre_protocol::markdown`]).
3//!
4//! Templates read `results.<query>`: `value` (first column of the first row), `first.<column>`,
5//! `rows` (at most `max_rows`), `row_count` (exact), `columns` and `sets[n]`. With neither
6//! `text:` nor `file:`, [`default_text`] writes one block per query.
7
8use std::fs::File;
9use std::path::Path;
10use std::sync::Arc;
11
12use arrow::ipc::reader::FileReader;
13use dre_protocol::markdown::escape;
14use minijinja::value::{Enumerator, Object, ObjectRepr, Value};
15
16use crate::dates::Calendar;
17use crate::numbers::Locale;
18
19/// A message output's own options (besides the shared output keys).
20pub const MESSAGE_KEYS: &[&str] = &["text", "file", "title", "max_rows"];
21/// How many rows of each query `results.<query>.rows` holds by default.
22pub const DEFAULT_MAX_ROWS: u64 = 1_000;
23/// How many rows the default template lists before `+ N more`.
24const DEFAULT_LIST_ROWS: usize = 10;
25
26/// One query's result as message templates see it.
27#[derive(Debug)]
28pub struct QueryResult {
29    pub columns: Arc<Vec<String>>,
30    /// At most the cap's rows.
31    pub rows: Vec<Vec<Value>>,
32    /// Every row the query returned.
33    pub row_count: u64,
34}
35
36impl QueryResult {
37    /// Read up to `max_rows` rows of a spooled result set.
38    pub fn read(spool: &Path, row_count: u64, max_rows: u64, cal: Calendar) -> Result<QueryResult, String> {
39        let reader = FileReader::try_new(File::open(spool).map_err(|e| e.to_string())?, None)
40            .map_err(|e| e.to_string())?;
41        let columns = Arc::new(
42            reader
43                .schema()
44                .fields()
45                .iter()
46                .map(|f| f.name().clone())
47                .collect(),
48        );
49        let mut rows = Vec::new();
50        for b in reader {
51            if rows.len() as u64 >= max_rows {
52                break;
53            }
54            let b = b.map_err(|e| e.to_string())?;
55            let left = (max_rows - rows.len() as u64) as usize;
56            let b = if b.num_rows() > left { b.slice(0, left) } else { b };
57            rows.extend(crate::values::typed_rows(&b, cal));
58        }
59        Ok(QueryResult {
60            columns,
61            rows,
62            row_count,
63        })
64    }
65
66    pub fn capped(&self) -> bool {
67        (self.rows.len() as u64) < self.row_count
68    }
69
70    /// The template value: an object with `value`, `first`, `rows`, `row_count`, `columns` and
71    /// `sets`.
72    pub fn value(self: Arc<Self>) -> Value {
73        Value::from_object(ResultValue(self))
74    }
75
76    fn row(&self, i: usize) -> Option<Value> {
77        self.rows
78            .get(i)
79            .map(|r| crate::render::Row::value(self.columns.clone(), r.clone()))
80    }
81}
82
83#[derive(Debug)]
84struct ResultValue(Arc<QueryResult>);
85
86impl Object for ResultValue {
87    fn repr(self: &Arc<Self>) -> ObjectRepr {
88        ObjectRepr::Map
89    }
90
91    fn get_value(self: &Arc<Self>, key: &Value) -> Option<Value> {
92        let r = &self.0;
93        Some(match key.as_str()? {
94            "value" => r
95                .rows
96                .first()
97                .and_then(|row| row.first().cloned())
98                .unwrap_or(Value::from(())),
99            "first" => r.row(0).unwrap_or(Value::from(())),
100            "rows" => Value::from((0..r.rows.len()).filter_map(|i| r.row(i)).collect::<Vec<_>>()),
101            "row_count" => Value::from(r.row_count),
102            "columns" => Value::from(r.columns.as_ref().clone()),
103            // One .sql file makes one result set (its last statement's), so `sets` holds it once.
104            "sets" => Value::from(vec![Value::from_object(ResultValue(r.clone()))]),
105            _ => return None,
106        })
107    }
108
109    fn enumerate(self: &Arc<Self>) -> Enumerator {
110        Enumerator::Str(&["value", "first", "rows", "row_count", "columns", "sets"])
111    }
112}
113
114/// A value as the default template shows it: numbers through `number` (two decimals unless
115/// whole), everything else as text; escaped for Markdown.
116fn shown(v: &Value, locale: Locale) -> String {
117    if v.is_none() {
118        return String::new();
119    }
120    let n = match v.as_i64() {
121        Some(i) => Some(i as f64),
122        None if v.kind() == minijinja::value::ValueKind::Number => f64::try_from(v.clone()).ok(),
123        None => None,
124    };
125    let text = match n {
126        Some(x) if x.fract() == 0.0 => locale.number(x, 0),
127        Some(x) => locale.number(x, 2),
128        None => v.to_string(),
129    };
130    escape(&text)
131}
132
133/// The default message: one block per query. A single value is `label: value`; one row is
134/// `label: value` lines; several rows are a list of at most ten, then `+ N more`.
135pub fn default_text(results: &[(String, Arc<QueryResult>)], locale: Locale) -> String {
136    let mut blocks = Vec::new();
137    for (query, r) in results {
138        let label = |c: &str| escape(c);
139        let block = match (r.row_count, r.columns.len()) {
140            (0, _) => format!("{}: no rows", label(query)),
141            (1, 1) => format!("{}: {}", label(&r.columns[0]), shown(&r.rows[0][0], locale)),
142            (1, _) => r
143                .columns
144                .iter()
145                .zip(&r.rows[0])
146                .map(|(c, v)| format!("{}: {}", label(c), shown(v, locale)))
147                .collect::<Vec<_>>()
148                .join("\n"),
149            (n, _) => {
150                let mut lines = vec![format!("**{}**", label(query))];
151                for row in r.rows.iter().take(DEFAULT_LIST_ROWS) {
152                    let cells: Vec<String> = row.iter().map(|v| shown(v, locale)).collect();
153                    lines.push(format!("- {}", cells.join(", ")));
154                }
155                let listed = r.rows.len().min(DEFAULT_LIST_ROWS) as u64;
156                if n > listed {
157                    lines.push(format!("+ {} more", crate::run::thousands(n - listed)));
158                }
159                lines.join("\n")
160            }
161        };
162        blocks.push(block);
163    }
164    blocks.join("\n\n")
165}
166
167/// What a message output writes to its `.md` file.
168pub fn file_body(title: &str, text: &str) -> String {
169    format!("# {title}\n\n{text}\n")
170}
171
172/// One line for the log: the title and the start of the text, as plain text.
173pub fn excerpt(title: &str, text: &str, max: usize) -> String {
174    let plain = dre_protocol::markdown::to_plain(text).replace('\n', " · ");
175    let line = format!("{title} — {plain}");
176    if line.chars().count() <= max {
177        line
178    } else {
179        format!(
180            "{}…",
181            line.chars().take(max.saturating_sub(1)).collect::<String>()
182        )
183    }
184}