1use std::time::Duration;
60
61use rich::{Console, ConsoleOptions, Justify, Renderable, Segment};
62use rich_ext::table::virtualized::{DEFAULT_SAMPLE, MAX_HEIGHT};
63use rich_ext::table::{Column, VirtualRows, VirtualTable};
64
65use crate::window::RowWindow;
66use crate::{DataError, DataType, Row, RowSource, Rows, Schema, Value};
67
68pub const DEFAULT_LIMIT: usize = 1_000;
70
71pub fn alignment(data_type: &DataType) -> Justify {
74 match data_type {
75 DataType::Integer | DataType::Float | DataType::Decimal { .. } => Justify::Right,
76 DataType::Boolean => Justify::Center,
77 _ => Justify::Left,
78 }
79}
80
81pub fn typed_columns(columns: &[String], schema: Option<&Schema>, sample: &[Row]) -> Vec<Column> {
85 columns
86 .iter()
87 .enumerate()
88 .map(|(index, name)| {
89 let field = schema.and_then(|s| s.field(name).or_else(|| s.fields().get(index)));
90 let justify = match field.map(|f| f.data_type()) {
91 Some(DataType::Any | DataType::Unknown) | None => {
92 let mut cells = sample
93 .iter()
94 .filter_map(|row| row.get(index))
95 .filter(|v| !matches!(v, Value::Null))
96 .peekable();
97 let numeric = cells.peek().is_some()
98 && cells.all(|v| matches!(v, Value::Int(_) | Value::Float(_)));
99 if numeric {
100 Justify::Right
101 } else {
102 Justify::Left
103 }
104 }
105 Some(data_type) => alignment(data_type),
106 };
107 Column::new(name.clone()).justify(justify)
108 })
109 .collect()
110}
111
112#[derive(Clone, Debug)]
115pub struct ResultSet<S = Rows> {
116 columns: Vec<String>,
117 schema: Option<Schema>,
118 source: S,
119 offset: usize,
120 limit: usize,
121 elapsed: Option<Duration>,
122 null: String,
123 title: Option<String>,
124}
125
126impl ResultSet<Rows> {
127 pub fn new(rows: Rows) -> Self {
129 let columns = rows.columns().to_vec();
130 let schema = rows.schema().cloned();
131 ResultSet::from_parts(columns, schema, rows)
132 }
133}
134
135impl ResultSet<RowWindow> {
136 pub fn read(source: impl RowSource, offset: usize, limit: usize) -> Result<Self, DataError> {
139 Ok(ResultSet::window(RowWindow::read(source, offset, limit)?))
140 }
141
142 pub fn window(window: RowWindow) -> Self {
144 let columns = window.columns().to_vec();
145 let schema = window.schema().cloned();
146 let (offset, limit) = (window.start(), window.rows().len());
147 ResultSet::from_parts(columns, schema, window)
148 .offset(offset)
149 .limit(limit)
150 }
151}
152
153impl<S: VirtualRows> ResultSet<S> {
154 pub fn from_parts(
157 columns: impl IntoIterator<Item = impl Into<String>>,
158 schema: Option<Schema>,
159 source: S,
160 ) -> Self {
161 ResultSet {
162 columns: columns.into_iter().map(Into::into).collect(),
163 schema,
164 source,
165 offset: 0,
166 limit: DEFAULT_LIMIT,
167 elapsed: None,
168 null: "NULL".to_string(),
169 title: None,
170 }
171 }
172
173 pub fn offset(mut self, offset: usize) -> Self {
175 self.offset = offset;
176 self
177 }
178
179 pub fn limit(mut self, rows: usize) -> Self {
182 self.limit = rows.min(MAX_HEIGHT);
183 self
184 }
185
186 pub fn elapsed(mut self, elapsed: Duration) -> Self {
188 self.elapsed = Some(elapsed);
189 self
190 }
191
192 pub fn null_marker(mut self, marker: impl Into<String>) -> Self {
194 self.null = marker.into();
195 self
196 }
197
198 pub fn title(mut self, title: impl Into<String>) -> Self {
200 self.title = Some(title.into());
201 self
202 }
203
204 pub fn columns(&self) -> &[String] {
206 &self.columns
207 }
208
209 pub fn schema(&self) -> Option<&Schema> {
211 self.schema.as_ref()
212 }
213
214 pub fn source(&self) -> &S {
216 &self.source
217 }
218
219 pub fn row_count(&self) -> Option<usize> {
221 self.source.row_count()
222 }
223
224 pub fn footer(&self, ascii: bool) -> String {
228 let count = match self.row_count() {
229 Some(count) => rows(count, ""),
230 None => {
231 let page = self.table(true).page();
232 match page.total {
233 Some(total) => rows(total, ""),
234 None => rows(page.range().end + usize::from(page.more), "+"),
235 }
236 }
237 };
238 match self.elapsed {
239 Some(elapsed) => format!(
240 "({count}, {})",
241 rich_ext::format::duration_with(elapsed, ascii)
242 ),
243 None => format!("({count})"),
244 }
245 }
246
247 pub fn is_windowed(&self) -> bool {
249 match self.row_count() {
250 Some(count) => (self.offset > 0 && count > 0) || count > self.limit,
251 None => true,
252 }
253 }
254
255 pub fn to_virtual_table(&self) -> VirtualTable<&S> {
258 self.table(self.is_windowed())
259 }
260
261 fn table(&self, position: bool) -> VirtualTable<&S> {
262 let sample = self.source.rows(0, DEFAULT_SAMPLE);
263 let mut table = VirtualTable::new(
264 typed_columns(&self.columns, self.schema.as_ref(), &sample),
265 &self.source,
266 )
267 .offset(self.offset)
268 .height(self.limit)
269 .null_marker(self.null.clone())
270 .fit_window(true)
271 .show_position(position);
272 if let Some(title) = &self.title {
273 table = table.title(title.clone());
274 }
275 table
276 }
277}
278
279fn rows(count: usize, suffix: &str) -> String {
281 let number = rich_ext::format::number(i64::try_from(count).unwrap_or(i64::MAX));
282 let noun = if count == 1 && suffix.is_empty() {
283 "row"
284 } else {
285 "rows"
286 };
287 format!("{number}{suffix} {noun}")
288}
289
290impl<S: VirtualRows> Renderable for ResultSet<S> {
291 fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
292 self.to_virtual_table()
293 .footnote(self.footer(console.ascii_only()))
294 .rich_render(console, options)
295 }
296
297 fn measure(&self, console: &Console, options: &ConsoleOptions) -> rich::measure::Measurement {
298 self.to_virtual_table()
299 .footnote(self.footer(console.ascii_only()))
300 .measure(console, options)
301 }
302}
303
304#[cfg(test)]
305mod tests {
306 use rich_ext::table::FnRows;
307
308 use super::*;
309 use crate::Field;
310
311 fn render(result: &impl Renderable, width: usize) -> String {
312 Console::builder()
313 .width(width)
314 .color_system(None)
315 .build()
316 .render_export(result)
317 }
318
319 #[test]
320 fn counts_read_as_words() {
321 let one = ResultSet::new({
322 let mut rows = Rows::new(["x"]);
323 rows.push([Value::Int(1)]);
324 rows
325 });
326 assert_eq!(one.footer(false), "(1 row)");
327 assert_eq!(ResultSet::new(Rows::new(["x"])).footer(false), "(0 rows)");
328 let many = ResultSet::from_parts(["x"], None, FnRows::new(Some(1_234_567), |_| None));
329 assert_eq!(many.footer(false), "(1,234,567 rows)");
330 }
331
332 #[test]
333 fn columns_without_types_align_by_their_cells() {
334 let sample = vec![
335 vec![Value::Int(1), "a".into(), Value::Null, Value::Null],
336 vec![Value::Float(2.5), "b".into(), Value::Null, Value::Int(3)],
337 ];
338 let names: Vec<String> = ["n", "s", "empty", "mixed"].map(String::from).to_vec();
339 let columns = typed_columns(&names, None, &sample);
340 let justify: Vec<Justify> = columns.iter().map(|c| c.column_options().justify).collect();
341 assert_eq!(
342 justify,
343 [Justify::Right, Justify::Left, Justify::Left, Justify::Right]
344 );
345 }
346
347 #[test]
348 fn schema_types_win_over_cells_and_match_by_name() {
349 let schema = Schema::new([
350 Field::new("b", DataType::Integer),
351 Field::new("a", DataType::Date),
352 ]);
353 let sample = vec![vec![Value::Int(1), "x".into()]];
354 let names: Vec<String> = ["a", "b"].map(String::from).to_vec();
355 let columns = typed_columns(&names, Some(&schema), &sample);
356 assert_eq!(columns[0].column_options().justify, Justify::Left);
357 assert_eq!(columns[1].column_options().justify, Justify::Right);
358 }
359
360 #[test]
361 fn a_large_result_shows_its_window_and_count() {
362 let source = FnRows::new(Some(1_000_000), |i| Some(vec![Value::from(i)]));
363 let result = ResultSet::from_parts(["i"], None, source)
364 .offset(10)
365 .limit(2);
366 let out = render(&result, 40);
367 assert!(
368 out.ends_with("│ 10 │\n│ 11 │\n└────┘\nrows 11–12 of 1,000,000\n(1,000,000 rows)\n"),
369 "{out}"
370 );
371 }
372
373 #[test]
374 fn an_uncounted_window_says_at_least() {
375 let mut rows = Rows::new(["n"]);
376 for n in 0..50 {
377 rows.push([Value::Int(n)]);
378 }
379 let window = RowWindow::reader()
380 .len(5)
381 .count(false)
382 .read(rows.into_source())
383 .unwrap();
384 let result = ResultSet::window(window);
385 assert_eq!(result.footer(false), "(6+ rows)");
386 let out = render(&result, 30);
387 assert!(out.ends_with("rows 1–5 of 6+\n(6+ rows)\n"), "{out}");
388 }
389
390 #[test]
391 fn the_measurement_covers_the_footer() {
392 let mut rows = Rows::new(["n"]);
393 rows.push([Value::Int(1)]);
394 let result = ResultSet::new(rows).elapsed(Duration::from_millis(12));
395 let console = Console::builder().width(40).build();
396 let measured = rich::measure::Measurement::get(&console, &console.options(), &result);
397 assert_eq!(measured.maximum, "(1 row, 12ms)".len());
398 }
399
400 #[test]
401 fn a_short_result_has_no_position_line() {
402 let mut rows = Rows::new(["a"]);
403 rows.push([Value::from("")]);
404 let out = render(&ResultSet::new(rows).title("q"), 30);
405 assert!(!out.contains("rows 1"), "{out}");
406 assert!(out.ends_with("(1 row)\n"), "{out}");
407 }
408}