Skip to main content

hematite/sql/
result.rs

1//! SQL result set and row interface
2
3use crate::error::{HematiteError, Result};
4use crate::query::{
5    DateTimeValue, DateValue, DecimalValue, IntervalDaySecondValue, IntervalYearMonthValue,
6    QueryResult, TimeValue, TimeWithTimeZoneValue, Value,
7};
8use std::collections::HashMap;
9
10#[derive(Debug, Clone)]
11pub struct ResultSet {
12    pub columns: Vec<String>,
13    pub rows: Vec<Row>,
14    column_index: HashMap<String, usize>,
15}
16
17impl ResultSet {
18    pub fn new(columns: Vec<String>, rows: Vec<Vec<Value>>) -> Self {
19        let mut column_index = HashMap::new();
20        for (i, col) in columns.iter().enumerate() {
21            column_index.insert(col.clone(), i);
22        }
23
24        let rows = rows.into_iter().map(Row::new).collect();
25
26        Self {
27            columns,
28            rows,
29            column_index,
30        }
31    }
32
33    pub fn len(&self) -> usize {
34        self.rows.len()
35    }
36
37    pub fn is_empty(&self) -> bool {
38        self.rows.is_empty()
39    }
40
41    pub fn get_row(&self, index: usize) -> Option<&Row> {
42        self.rows.get(index)
43    }
44
45    pub fn iter(&'_ self) -> std::slice::Iter<'_, Row> {
46        self.rows.iter()
47    }
48
49    pub fn column_count(&self) -> usize {
50        self.columns.len()
51    }
52
53    pub fn get_column_index(&self, column_name: &str) -> Option<usize> {
54        self.column_index.get(column_name).copied()
55    }
56
57    pub fn to_structs<T: FromRow>(&self) -> Result<Vec<T>> {
58        self.rows.iter().map(T::from_row).collect()
59    }
60
61    pub fn render_ascii_table(&self) -> String {
62        if self.columns.is_empty() {
63            return format!("{} row(s)", self.len());
64        }
65
66        let mut widths = self
67            .columns
68            .iter()
69            .map(|column| display_width(column))
70            .collect::<Vec<_>>();
71
72        let rendered_rows = self
73            .rows
74            .iter()
75            .map(|row| {
76                row.values
77                    .iter()
78                    .map(render_value_for_table)
79                    .collect::<Vec<_>>()
80            })
81            .collect::<Vec<_>>();
82
83        for row in &rendered_rows {
84            for (index, value) in row.iter().enumerate() {
85                if index >= widths.len() {
86                    widths.push(display_width(value));
87                } else {
88                    widths[index] = widths[index].max(display_width(value));
89                }
90            }
91        }
92
93        let border = ascii_border(&widths);
94        let mut lines = Vec::new();
95        lines.push(border.clone());
96        lines.push(ascii_row(&self.columns, &widths));
97        lines.push(border.clone());
98        for row in &rendered_rows {
99            lines.push(ascii_row(row, &widths));
100        }
101        lines.push(border);
102        lines.push(format!("{} row(s)", self.len()));
103        lines.join("\n")
104    }
105}
106
107impl IntoIterator for ResultSet {
108    type Item = Row;
109    type IntoIter = std::vec::IntoIter<Row>;
110
111    fn into_iter(self) -> Self::IntoIter {
112        self.rows.into_iter()
113    }
114}
115
116#[derive(Debug, Clone)]
117pub struct Row {
118    pub values: Vec<Value>,
119}
120
121impl Row {
122    pub fn new(values: Vec<Value>) -> Self {
123        Self { values }
124    }
125
126    pub fn len(&self) -> usize {
127        self.values.len()
128    }
129
130    pub fn is_empty(&self) -> bool {
131        self.values.is_empty()
132    }
133
134    pub fn get(&self, index: usize) -> Option<&Value> {
135        self.values.get(index)
136    }
137
138    pub fn get_by_name(
139        &self,
140        column_name: &str,
141        column_index: &HashMap<String, usize>,
142    ) -> Option<&Value> {
143        if let Some(&idx) = column_index.get(column_name) {
144            self.get(idx)
145        } else {
146            None
147        }
148    }
149
150    pub fn to_struct<T: FromRow>(&self) -> Result<T> {
151        T::from_row(self)
152    }
153
154    pub fn get_int(&self, index: usize) -> Result<i32> {
155        match self.get(index) {
156            Some(Value::Integer(i)) => Ok(*i),
157            Some(value) => Err(HematiteError::ParseError(format!(
158                "Expected INT, found {:?}",
159                value
160            ))),
161            None => Err(HematiteError::ParseError(
162                "Column index out of bounds".to_string(),
163            )),
164        }
165    }
166
167    pub fn get_string(&self, index: usize) -> Result<String> {
168        match self.get(index) {
169            Some(Value::Text(s)) => Ok(s.clone()),
170            Some(Value::Enum(s)) => Ok(s.clone()),
171            Some(Value::Decimal(value)) => Ok(value.to_string()),
172            Some(Value::Date(value)) => Ok(value.to_string()),
173            Some(Value::Time(value)) => Ok(value.to_string()),
174            Some(Value::DateTime(value)) => Ok(value.to_string()),
175            Some(Value::TimeWithTimeZone(value)) => Ok(value.to_string()),
176            Some(Value::IntervalYearMonth(value)) => Ok(value.to_string()),
177            Some(Value::IntervalDaySecond(value)) => Ok(value.to_string()),
178            Some(value) => Err(HematiteError::ParseError(format!(
179                "Expected TEXT, found {:?}",
180                value
181            ))),
182            None => Err(HematiteError::ParseError(
183                "Column index out of bounds".to_string(),
184            )),
185        }
186    }
187
188    pub fn get_bool(&self, index: usize) -> Result<bool> {
189        match self.get(index) {
190            Some(Value::Boolean(b)) => Ok(*b),
191            Some(value) => Err(HematiteError::ParseError(format!(
192                "Expected BOOLEAN, found {:?}",
193                value
194            ))),
195            None => Err(HematiteError::ParseError(
196                "Column index out of bounds".to_string(),
197            )),
198        }
199    }
200
201    pub fn get_float(&self, index: usize) -> Result<f64> {
202        match self.get(index) {
203            Some(Value::Float32(f)) => Ok(*f as f64),
204            Some(Value::Float(f)) => Ok(*f),
205            Some(Value::Integer(i)) => Ok(*i as f64), // Allow integer to float conversion
206            Some(Value::UInteger(i)) => Ok(*i as f64),
207            Some(Value::BigInt(i)) => Ok(*i as f64),
208            Some(Value::UBigInt(i)) => Ok(*i as f64),
209            Some(Value::Int128(i)) => Ok(*i as f64),
210            Some(Value::UInt128(i)) => Ok(*i as f64),
211            Some(value) => Err(HematiteError::ParseError(format!(
212                "Expected FLOAT, found {:?}",
213                value
214            ))),
215            None => Err(HematiteError::ParseError(
216                "Column index out of bounds".to_string(),
217            )),
218        }
219    }
220
221    pub fn is_null(&self, index: usize) -> bool {
222        matches!(self.get(index), Some(Value::Null))
223    }
224
225    pub fn get_bigint(&self, index: usize) -> Result<i64> {
226        match self.get(index) {
227            Some(Value::BigInt(i)) => Ok(*i),
228            Some(Value::Integer(i)) => Ok(*i as i64),
229            Some(Value::UInteger(i)) => Ok(*i as i64),
230            Some(value) => Err(HematiteError::ParseError(format!(
231                "Expected INT64, found {:?}",
232                value
233            ))),
234            None => Err(HematiteError::ParseError(
235                "Column index out of bounds".to_string(),
236            )),
237        }
238    }
239
240    pub fn get_int128(&self, index: usize) -> Result<i128> {
241        match self.get(index) {
242            Some(Value::Int128(i)) => Ok(*i),
243            Some(Value::BigInt(i)) => Ok(*i as i128),
244            Some(Value::Integer(i)) => Ok(*i as i128),
245            Some(Value::UInteger(i)) => Ok(*i as i128),
246            Some(Value::UBigInt(i)) => Ok(*i as i128),
247            Some(value) => Err(HematiteError::ParseError(format!(
248                "Expected INT128, found {:?}",
249                value
250            ))),
251            None => Err(HematiteError::ParseError(
252                "Column index out of bounds".to_string(),
253            )),
254        }
255    }
256
257    pub fn get_uint(&self, index: usize) -> Result<u32> {
258        match self.get(index) {
259            Some(Value::UInteger(i)) => Ok(*i),
260            Some(Value::Integer(i)) if *i >= 0 => Ok(*i as u32),
261            Some(value) => Err(HematiteError::ParseError(format!(
262                "Expected UINT, found {:?}",
263                value
264            ))),
265            None => Err(HematiteError::ParseError(
266                "Column index out of bounds".to_string(),
267            )),
268        }
269    }
270
271    pub fn get_uint64(&self, index: usize) -> Result<u64> {
272        match self.get(index) {
273            Some(Value::UBigInt(i)) => Ok(*i),
274            Some(Value::UInteger(i)) => Ok(*i as u64),
275            Some(Value::Integer(i)) if *i >= 0 => Ok(*i as u64),
276            Some(value) => Err(HematiteError::ParseError(format!(
277                "Expected UINT64, found {:?}",
278                value
279            ))),
280            None => Err(HematiteError::ParseError(
281                "Column index out of bounds".to_string(),
282            )),
283        }
284    }
285
286    pub fn get_uint128(&self, index: usize) -> Result<u128> {
287        match self.get(index) {
288            Some(Value::UInt128(i)) => Ok(*i),
289            Some(Value::UBigInt(i)) => Ok(*i as u128),
290            Some(Value::UInteger(i)) => Ok(*i as u128),
291            Some(Value::Integer(i)) if *i >= 0 => Ok(*i as u128),
292            Some(value) => Err(HematiteError::ParseError(format!(
293                "Expected UINT128, found {:?}",
294                value
295            ))),
296            None => Err(HematiteError::ParseError(
297                "Column index out of bounds".to_string(),
298            )),
299        }
300    }
301
302    pub fn get_decimal(&self, index: usize) -> Result<DecimalValue> {
303        match self.get(index) {
304            Some(Value::Decimal(value)) => Ok(value.clone()),
305            Some(value) => Err(HematiteError::ParseError(format!(
306                "Expected DECIMAL, found {:?}",
307                value
308            ))),
309            None => Err(HematiteError::ParseError(
310                "Column index out of bounds".to_string(),
311            )),
312        }
313    }
314
315    pub fn get_blob(&self, index: usize) -> Result<Vec<u8>> {
316        match self.get(index) {
317            Some(Value::Blob(value)) => Ok(value.clone()),
318            Some(value) => Err(HematiteError::ParseError(format!(
319                "Expected BLOB, found {:?}",
320                value
321            ))),
322            None => Err(HematiteError::ParseError(
323                "Column index out of bounds".to_string(),
324            )),
325        }
326    }
327
328    pub fn get_date(&self, index: usize) -> Result<DateValue> {
329        match self.get(index) {
330            Some(Value::Date(value)) => Ok(*value),
331            Some(value) => Err(HematiteError::ParseError(format!(
332                "Expected DATE, found {:?}",
333                value
334            ))),
335            None => Err(HematiteError::ParseError(
336                "Column index out of bounds".to_string(),
337            )),
338        }
339    }
340
341    pub fn get_time(&self, index: usize) -> Result<TimeValue> {
342        match self.get(index) {
343            Some(Value::Time(value)) => Ok(*value),
344            Some(value) => Err(HematiteError::ParseError(format!(
345                "Expected TIME, found {:?}",
346                value
347            ))),
348            None => Err(HematiteError::ParseError(
349                "Column index out of bounds".to_string(),
350            )),
351        }
352    }
353
354    pub fn get_datetime(&self, index: usize) -> Result<DateTimeValue> {
355        match self.get(index) {
356            Some(Value::DateTime(value)) => Ok(*value),
357            Some(value) => Err(HematiteError::ParseError(format!(
358                "Expected DATETIME, found {:?}",
359                value
360            ))),
361            None => Err(HematiteError::ParseError(
362                "Column index out of bounds".to_string(),
363            )),
364        }
365    }
366
367    pub fn get_time_with_time_zone(&self, index: usize) -> Result<TimeWithTimeZoneValue> {
368        match self.get(index) {
369            Some(Value::TimeWithTimeZone(value)) => Ok(*value),
370            Some(value) => Err(HematiteError::ParseError(format!(
371                "Expected TIME WITH TIME ZONE, found {:?}",
372                value
373            ))),
374            None => Err(HematiteError::ParseError(
375                "Column index out of bounds".to_string(),
376            )),
377        }
378    }
379
380    pub fn get_interval_year_month(&self, index: usize) -> Result<IntervalYearMonthValue> {
381        match self.get(index) {
382            Some(Value::IntervalYearMonth(value)) => Ok(*value),
383            Some(value) => Err(HematiteError::ParseError(format!(
384                "Expected INTERVAL YEAR TO MONTH, found {:?}",
385                value
386            ))),
387            None => Err(HematiteError::ParseError(
388                "Column index out of bounds".to_string(),
389            )),
390        }
391    }
392
393    pub fn get_interval_day_second(&self, index: usize) -> Result<IntervalDaySecondValue> {
394        match self.get(index) {
395            Some(Value::IntervalDaySecond(value)) => Ok(*value),
396            Some(value) => Err(HematiteError::ParseError(format!(
397                "Expected INTERVAL DAY TO SECOND, found {:?}",
398                value
399            ))),
400            None => Err(HematiteError::ParseError(
401                "Column index out of bounds".to_string(),
402            )),
403        }
404    }
405}
406
407pub trait FromRow: Sized {
408    fn from_row(row: &Row) -> Result<Self>;
409}
410
411impl IntoIterator for Row {
412    type Item = Value;
413    type IntoIter = std::vec::IntoIter<Value>;
414
415    fn into_iter(self) -> Self::IntoIter {
416        self.values.into_iter()
417    }
418}
419
420#[derive(Debug, Clone)]
421pub struct StatementResult {
422    pub affected_rows: usize,
423    pub last_insert_id: Option<i32>,
424    pub message: String,
425}
426
427impl StatementResult {
428    pub fn new(affected_rows: usize, message: String) -> Self {
429        Self {
430            affected_rows,
431            last_insert_id: None,
432            message,
433        }
434    }
435
436    pub fn with_insert_id(affected_rows: usize, last_insert_id: i32, message: String) -> Self {
437        Self {
438            affected_rows,
439            last_insert_id: Some(last_insert_id),
440            message,
441        }
442    }
443}
444
445#[derive(Debug, Clone)]
446pub enum ExecutedStatement {
447    Query(ResultSet),
448    Statement(StatementResult),
449}
450
451impl ExecutedStatement {
452    pub(crate) fn from_query_result(query_result: QueryResult) -> Self {
453        if query_result.columns.is_empty() {
454            Self::Statement(StatementResult::new(
455                query_result.affected_rows,
456                "Ok".to_string(),
457            ))
458        } else {
459            Self::Query(ResultSet::new(query_result.columns, query_result.rows))
460        }
461    }
462
463    pub fn as_query(&self) -> Option<&ResultSet> {
464        match self {
465            Self::Query(result_set) => Some(result_set),
466            Self::Statement(_) => None,
467        }
468    }
469
470    pub fn as_statement(&self) -> Option<&StatementResult> {
471        match self {
472            Self::Query(_) => None,
473            Self::Statement(result) => Some(result),
474        }
475    }
476
477    pub fn render_ascii(&self) -> String {
478        match self {
479            Self::Query(result_set) => result_set.render_ascii_table(),
480            Self::Statement(result) => format!("{} ({})", result.message, result.affected_rows),
481        }
482    }
483}
484
485fn render_value_for_table(value: &Value) -> String {
486    match value {
487        Value::Null => "NULL".to_string(),
488        Value::Text(value) => sanitize_table_cell(value),
489        Value::Enum(value) => sanitize_table_cell(value),
490        Value::Blob(bytes) => sanitize_table_cell(&format!("0x{}", hex::encode(bytes))),
491        Value::Boolean(value) => value.to_string(),
492        Value::Integer(value) => value.to_string(),
493        Value::BigInt(value) => value.to_string(),
494        Value::Int128(value) => value.to_string(),
495        Value::UInteger(value) => value.to_string(),
496        Value::UBigInt(value) => value.to_string(),
497        Value::UInt128(value) => value.to_string(),
498        Value::Float32(value) => value.to_string(),
499        Value::Float(value) => value.to_string(),
500        Value::Decimal(value) => value.to_string(),
501        Value::Date(value) => value.to_string(),
502        Value::Time(value) => value.to_string(),
503        Value::DateTime(value) => value.to_string(),
504        Value::TimeWithTimeZone(value) => value.to_string(),
505        Value::IntervalYearMonth(value) => value.to_string(),
506        Value::IntervalDaySecond(value) => value.to_string(),
507    }
508}
509
510fn sanitize_table_cell(value: &str) -> String {
511    value.replace('\n', "\\n").replace('\r', "\\r")
512}
513
514fn display_width(value: &str) -> usize {
515    value.chars().count()
516}
517
518fn ascii_border(widths: &[usize]) -> String {
519    let mut border = String::new();
520    border.push('+');
521    for width in widths {
522        border.push_str(&"-".repeat(*width + 2));
523        border.push('+');
524    }
525    border
526}
527
528fn ascii_row(values: &[String], widths: &[usize]) -> String {
529    let mut row = String::new();
530    row.push('|');
531    for (index, width) in widths.iter().enumerate() {
532        let value = values.get(index).map(String::as_str).unwrap_or("");
533        row.push(' ');
534        row.push_str(value);
535        let padding = width.saturating_sub(display_width(value));
536        row.push_str(&" ".repeat(padding));
537        row.push(' ');
538        row.push('|');
539    }
540    row
541}
542
543mod hex {
544    pub fn encode(bytes: &[u8]) -> String {
545        let mut output = String::with_capacity(bytes.len() * 2);
546        for byte in bytes {
547            output.push(nibble_to_hex(byte >> 4));
548            output.push(nibble_to_hex(byte & 0x0f));
549        }
550        output
551    }
552
553    fn nibble_to_hex(value: u8) -> char {
554        match value {
555            0..=9 => (b'0' + value) as char,
556            10..=15 => (b'A' + (value - 10)) as char,
557            _ => unreachable!("hex nibble out of range"),
558        }
559    }
560}