1use 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), 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}