1use 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
19pub const MESSAGE_KEYS: &[&str] = &["text", "file", "title", "max_rows"];
21pub const DEFAULT_MAX_ROWS: u64 = 1_000;
23const DEFAULT_LIST_ROWS: usize = 10;
25
26#[derive(Debug)]
28pub struct QueryResult {
29 pub columns: Arc<Vec<String>>,
30 pub rows: Vec<Vec<Value>>,
32 pub row_count: u64,
34}
35
36impl QueryResult {
37 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 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 "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
114fn 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
133pub 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
167pub fn file_body(title: &str, text: &str) -> String {
169 format!("# {title}\n\n{text}\n")
170}
171
172pub 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}