1use anyhow::{bail, ensure, Context, Result};
31use chrono::{DateTime, Duration as ChronoDuration, FixedOffset, NaiveDate};
32use serde_json::{json, Map, Value};
33
34use super::arrow::{self, ArrowType, Cell};
35
36pub const MAX_COLUMNS: usize = 1024;
38pub const MAX_CELL_BYTES: usize = 16 * 1024;
40pub const MAX_VARIANT_DEPTH: usize = 64;
42pub const MAX_DECODE_WARNINGS: usize = 16;
44
45#[derive(Debug, Clone, Copy, PartialEq)]
47pub enum Kind {
48 Fixed,
49 Real,
50 Text,
51 Boolean,
52 Date,
53 Time,
54 TimestampNtz,
55 TimestampLtz,
56 TimestampTz,
57 Variant,
58 Binary,
59 Other,
60}
61
62impl Kind {
63 pub fn from_type(name: &str) -> Kind {
64 match name.to_ascii_lowercase().as_str() {
65 "fixed" | "number" | "decimal" | "numeric" | "int" | "integer" | "bigint" => {
66 Kind::Fixed
67 }
68 "real" | "float" | "double" => Kind::Real,
69 "text" | "varchar" | "string" | "char" => Kind::Text,
70 "boolean" => Kind::Boolean,
71 "date" => Kind::Date,
72 "time" => Kind::Time,
73 "timestamp_ntz" | "timestamp" | "datetime" => Kind::TimestampNtz,
74 "timestamp_ltz" => Kind::TimestampLtz,
75 "timestamp_tz" => Kind::TimestampTz,
76 "variant" | "object" | "array" => Kind::Variant,
77 "binary" => Kind::Binary,
78 _ => Kind::Other,
79 }
80 }
81}
82
83#[derive(Debug, Clone)]
85pub struct Column {
86 pub name: String,
87 pub type_name: String,
88 pub kind: Kind,
89 pub nullable: bool,
90 pub precision: Option<i64>,
91 pub scale: u32,
92}
93
94impl Column {
95 pub fn from_rowtype(v: &Value) -> Result<Column> {
96 let name = v
97 .get("name")
98 .and_then(Value::as_str)
99 .context("a rowtype column has no name")?
100 .to_string();
101 let type_name = v
102 .get("type")
103 .and_then(Value::as_str)
104 .with_context(|| format!("rowtype column {name:?} has no type"))?
105 .to_ascii_lowercase();
106 let scale = match v.get("scale") {
107 None | Some(Value::Null) => None,
108 Some(s) => Some(
109 s.as_u64()
110 .filter(|s| *s <= 38)
111 .with_context(|| format!("column {name:?} has scale {s}"))?,
112 ),
113 };
114 let kind = Kind::from_type(&type_name);
115 let default_scale = match kind {
117 Kind::Time | Kind::TimestampNtz | Kind::TimestampLtz | Kind::TimestampTz => 9,
118 _ => 0,
119 };
120 let scale = scale.unwrap_or(default_scale) as u32;
121 if matches!(
122 kind,
123 Kind::Time | Kind::TimestampNtz | Kind::TimestampLtz | Kind::TimestampTz
124 ) {
125 ensure!(scale <= 9, "column {name:?}: a time scale of {scale}");
126 }
127 Ok(Column {
128 name,
129 type_name,
130 kind,
131 nullable: v.get("nullable").and_then(Value::as_bool).unwrap_or(true),
132 precision: v.get("precision").and_then(Value::as_i64),
133 scale,
134 })
135 }
136
137 pub fn to_json(&self) -> Value {
138 json!({
139 "name": self.name,
140 "type": self.type_name,
141 "nullable": self.nullable,
142 "precision": self.precision,
143 "scale": self.scale,
144 })
145 }
146}
147
148#[derive(Default)]
150struct Notes {
151 warnings: Vec<Value>,
152 warning_count: usize,
153 truncated_cells: usize,
154}
155
156impl Notes {
157 fn warn(&mut self, column: &Column, reason: String) {
158 self.warning_count += 1;
159 if self.warnings.len() < MAX_DECODE_WARNINGS {
160 self.warnings.push(json!({
161 "column": column.name,
162 "type": column.type_name,
163 "reason": reason,
164 }));
165 }
166 }
167 fn text(&mut self, s: &str) -> Value {
168 if s.len() > MAX_CELL_BYTES {
169 self.truncated_cells += 1;
170 Value::String(crate::utils::truncate::truncate_with_suffix(
171 s,
172 MAX_CELL_BYTES,
173 "...",
174 ))
175 } else {
176 Value::String(s.to_string())
177 }
178 }
179}
180
181pub fn json_depth(text: &str) -> usize {
183 let (mut depth, mut max, mut in_string, mut escaped) = (0usize, 0usize, false, false);
184 for b in text.bytes() {
185 if in_string {
186 match (escaped, b) {
187 (true, _) => escaped = false,
188 (false, b'\\') => escaped = true,
189 (false, b'"') => in_string = false,
190 _ => {}
191 }
192 continue;
193 }
194 match b {
195 b'"' => in_string = true,
196 b'[' | b'{' => {
197 depth += 1;
198 max = max.max(depth);
199 }
200 b']' | b'}' => depth = depth.saturating_sub(1),
201 _ => {}
202 }
203 }
204 max
205}
206
207pub fn decimal_string(units: i128, scale: u32) -> String {
209 if scale == 0 {
210 return units.to_string();
211 }
212 let digits = units.unsigned_abs().to_string();
213 let scale = scale as usize;
214 let padded = if digits.len() <= scale {
215 format!("{}{}", "0".repeat(scale + 1 - digits.len()), digits)
216 } else {
217 digits
218 };
219 let (int, frac) = padded.split_at(padded.len() - scale);
220 format!("{}{int}.{frac}", if units < 0 { "-" } else { "" })
221}
222
223fn integer(units: i128) -> Value {
225 match i64::try_from(units) {
226 Ok(v) => json!(v),
227 Err(_) => Value::String(units.to_string()),
228 }
229}
230
231pub fn parse_seconds(text: &str) -> Result<i128> {
233 let (neg, body) = match text.strip_prefix('-') {
234 Some(rest) => (true, rest),
235 None => (false, text),
236 };
237 let (int, frac) = body.split_once('.').unwrap_or((body, ""));
238 ensure!(
239 !int.is_empty()
240 && int.bytes().all(|b| b.is_ascii_digit())
241 && frac.bytes().all(|b| b.is_ascii_digit())
242 && frac.len() <= 9
243 && int.len() <= 18,
244 "{text:?} is not seconds with up to nine fraction digits"
245 );
246 let int: i128 = int.parse()?;
247 let frac: i128 = if frac.is_empty() {
248 0
249 } else {
250 format!("{frac:0<9}").parse()?
251 };
252 let nanos = int * 1_000_000_000 + frac;
253 Ok(if neg { -nanos } else { nanos })
254}
255
256fn fraction(nanos: u32, scale: u32) -> String {
257 if scale == 0 {
258 return String::new();
259 }
260 format!(".{}", &format!("{nanos:09}")[..scale as usize])
261}
262
263fn split_nanos(nanos: i128) -> Result<(i64, u32)> {
264 let secs = i64::try_from(nanos.div_euclid(1_000_000_000)).context("time out of range")?;
265 Ok((secs, nanos.rem_euclid(1_000_000_000) as u32))
266}
267
268pub fn date_from_days(days: i64) -> Result<String> {
269 let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).context("epoch")?;
270 let date = epoch
271 .checked_add_signed(ChronoDuration::try_days(days).context("date out of range")?)
272 .context("date out of range")?;
273 Ok(date.format("%Y-%m-%d").to_string())
274}
275
276pub fn time_from_nanos(nanos: i128, scale: u32) -> Result<String> {
277 ensure!(
278 (0..86_400 * 1_000_000_000).contains(&nanos),
279 "a time of day outside 00:00:00..24:00:00"
280 );
281 let (secs, frac) = split_nanos(nanos)?;
282 Ok(format!(
283 "{:02}:{:02}:{:02}{}",
284 secs / 3600,
285 secs / 60 % 60,
286 secs % 60,
287 fraction(frac, scale)
288 ))
289}
290
291pub fn timestamp(
293 nanos: i128,
294 scale: u32,
295 kind: Kind,
296 offset_minutes: Option<i32>,
297) -> Result<String> {
298 let (secs, frac) = split_nanos(nanos)?;
299 let utc = DateTime::from_timestamp(secs, frac).context("timestamp out of range")?;
300 let base = |t: chrono::NaiveDateTime| {
301 format!("{}{}", t.format("%Y-%m-%dT%H:%M:%S"), fraction(frac, scale))
302 };
303 Ok(match kind {
304 Kind::TimestampNtz => base(utc.naive_utc()),
305 Kind::TimestampLtz => format!("{}Z", base(utc.naive_utc())),
306 _ => {
307 let minutes = offset_minutes.context("TIMESTAMP_TZ without an offset")?;
308 let offset = FixedOffset::east_opt(minutes * 60)
309 .with_context(|| format!("a zone offset of {minutes} minutes"))?;
310 let local = utc.with_timezone(&offset);
311 format!("{}{}", base(local.naive_local()), local.format("%:z"))
312 }
313 })
314}
315
316fn zone(encoded: i128) -> Result<i32> {
318 let minutes = i32::try_from(encoded - 1440).context("zone out of range")?;
319 ensure!(
320 (-1440..=1440).contains(&minutes),
321 "a zone offset of {minutes} minutes"
322 );
323 Ok(minutes)
324}
325
326fn variant(text: &str, notes: &mut Notes, column: &Column) -> Value {
327 if json_depth(text) > MAX_VARIANT_DEPTH {
328 notes.warn(
329 column,
330 format!("nests deeper than {MAX_VARIANT_DEPTH}; given as text"),
331 );
332 return notes.text(text);
333 }
334 match serde_json::from_str::<Value>(text) {
335 Ok(v) if text.len() <= MAX_CELL_BYTES => v,
336 Ok(_) => notes.text(text),
337 Err(e) => {
338 notes.warn(column, format!("not JSON: {e}"));
339 notes.text(text)
340 }
341 }
342}
343
344fn json_cell(raw: &Value, column: &Column, notes: &mut Notes) -> Value {
346 let text = match raw {
347 Value::Null => return Value::Null,
348 Value::String(s) => s.as_str(),
349 other => {
350 notes.warn(column, "a JSON rowset cell that is not a string".into());
351 return other.clone();
352 }
353 };
354 let decoded: Result<Value> = (|| {
355 Ok(match column.kind {
356 Kind::Text | Kind::Other => notes.text(text),
357 Kind::Variant => variant(text, notes, column),
358 Kind::Fixed => {
359 if column.scale == 0 {
360 let v: i128 = text
361 .parse()
362 .with_context(|| format!("{text:?} is not a whole number"))?;
363 integer(v)
364 } else {
365 let (int, frac) = text.split_once('.').unwrap_or((text, ""));
366 let int = int.strip_prefix('-').unwrap_or(int);
367 ensure!(
368 !int.is_empty()
369 && int.bytes().all(|b| b.is_ascii_digit())
370 && frac.bytes().all(|b| b.is_ascii_digit())
371 && frac.len() <= column.scale as usize,
372 "{text:?} is not a number of scale {}",
373 column.scale
374 );
375 Value::String(text.to_string())
376 }
377 }
378 Kind::Real => {
379 let v: f64 = text
380 .parse()
381 .with_context(|| format!("{text:?} is not a number"))?;
382 if v.is_finite() {
383 json!(v)
384 } else {
385 Value::String(text.to_string())
386 }
387 }
388 Kind::Boolean => match text {
389 "1" | "TRUE" | "true" | "True" => Value::Bool(true),
390 "0" | "FALSE" | "false" | "False" => Value::Bool(false),
391 _ => bail!("{text:?} is not a boolean"),
392 },
393 Kind::Date => Value::String(date_from_days(
394 text.parse()
395 .with_context(|| format!("{text:?} is not a day count"))?,
396 )?),
397 Kind::Time => Value::String(time_from_nanos(parse_seconds(text)?, column.scale)?),
398 Kind::TimestampNtz | Kind::TimestampLtz => Value::String(timestamp(
399 parse_seconds(text)?,
400 column.scale,
401 column.kind,
402 None,
403 )?),
404 Kind::TimestampTz => {
405 let (value, tz) = text
406 .split_once(' ')
407 .with_context(|| format!("{text:?} carries no zone"))?;
408 let tz: i128 = tz.parse().with_context(|| format!("zone {tz:?}"))?;
409 Value::String(timestamp(
410 parse_seconds(value)?,
411 column.scale,
412 column.kind,
413 Some(zone(tz)?),
414 )?)
415 }
416 Kind::Binary => {
417 ensure!(
418 text.len().is_multiple_of(2) && text.bytes().all(|b| b.is_ascii_hexdigit()),
419 "{} characters that are not hex",
420 text.len()
421 );
422 json!({"binary_length": text.len() / 2})
423 }
424 })
425 })();
426 match decoded {
427 Ok(v) => v,
428 Err(e) => {
429 notes.warn(column, format!("{e:#}"));
430 notes.text(text)
431 }
432 }
433}
434
435fn scaled(units: i128, scale: u32) -> Result<i128> {
436 let factor = 10i128
437 .checked_pow(9u32.checked_sub(scale).context("a scale above 9")?)
438 .context("scale")?;
439 units.checked_mul(factor).context("time out of range")
440}
441
442fn raw_arrow(cell: &Cell) -> Value {
444 match cell {
445 Cell::Null => Value::Null,
446 Cell::Int(v) => integer(*v),
447 Cell::Float(f) => json!(f),
448 Cell::Bool(b) => json!(b),
449 Cell::Text(s) => json!(s),
450 Cell::Bytes(n) => json!({ "binary_length": n }),
451 Cell::Struct(c) => Value::Array(c.iter().map(raw_arrow).collect()),
452 }
453}
454
455fn arrow_cell(cell: &Cell, field: &arrow::Field, column: &Column, notes: &mut Notes) -> Value {
457 if *cell == Cell::Null {
458 return Value::Null;
459 }
460 if let Cell::Struct(parts) = cell {
463 if !parts.is_empty() && parts.iter().all(|c| *c == Cell::Null) {
464 return Value::Null;
465 }
466 }
467 let decoded: Result<Value> = (|| {
468 Ok(match (column.kind, cell, &field.data_type) {
469 (Kind::Fixed, Cell::Int(v), ArrowType::Decimal { scale, .. }) => {
470 ensure!(*scale >= 0, "a negative decimal scale");
471 if *scale == 0 {
472 integer(*v)
473 } else {
474 Value::String(decimal_string(*v, *scale as u32))
475 }
476 }
477 (Kind::Fixed, Cell::Int(v), ArrowType::Int { .. }) => {
478 if column.scale == 0 {
479 integer(*v)
480 } else {
481 Value::String(decimal_string(*v, column.scale))
482 }
483 }
484 (Kind::Fixed | Kind::Real, Cell::Float(f), _) => {
485 if f.is_finite() {
486 json!(f)
487 } else {
488 Value::String(f.to_string())
489 }
490 }
491 (Kind::Real, Cell::Int(v), ArrowType::Int { .. }) => integer(*v),
492 (Kind::Text | Kind::Other, Cell::Text(s), _) => notes.text(s),
493 (Kind::Variant, Cell::Text(s), _) => variant(s, notes, column),
494 (Kind::Boolean, Cell::Bool(b), _) => Value::Bool(*b),
495 (Kind::Binary, Cell::Bytes(n), _) => json!({ "binary_length": n }),
496 (Kind::Date, Cell::Int(v), ArrowType::Date { day: true }) => {
497 Value::String(date_from_days(i64::try_from(*v)?)?)
498 }
499 (Kind::Date, Cell::Int(v), ArrowType::Date { day: false }) => {
500 Value::String(date_from_days(i64::try_from(v.div_euclid(86_400_000))?)?)
501 }
502 (Kind::Time, Cell::Int(v), ArrowType::Int { .. }) => {
503 Value::String(time_from_nanos(scaled(*v, column.scale)?, column.scale)?)
504 }
505 (Kind::Time, Cell::Int(v), ArrowType::Time { scale, .. }) => {
506 Value::String(time_from_nanos(scaled(*v, *scale)?, column.scale)?)
507 }
508 (
509 Kind::TimestampNtz | Kind::TimestampLtz | Kind::TimestampTz,
510 Cell::Int(v),
511 ArrowType::Int { .. },
512 ) => {
513 ensure!(
514 column.kind != Kind::TimestampTz,
515 "TIMESTAMP_TZ as a bare integer carries no zone"
516 );
517 Value::String(timestamp(
518 scaled(*v, column.scale)?,
519 column.scale,
520 column.kind,
521 None,
522 )?)
523 }
524 (
525 Kind::TimestampNtz | Kind::TimestampLtz,
526 Cell::Int(v),
527 ArrowType::Timestamp { scale, .. },
528 ) => Value::String(timestamp(
529 scaled(*v, *scale)?,
530 column.scale,
531 column.kind,
532 None,
533 )?),
534 (
535 Kind::TimestampNtz | Kind::TimestampLtz | Kind::TimestampTz,
536 Cell::Struct(parts),
537 ArrowType::Struct,
538 ) => {
539 let get = |name: &str| -> Option<i128> {
540 field
541 .children
542 .iter()
543 .position(|c| c.name == name)
544 .and_then(|i| match parts.get(i) {
545 Some(Cell::Int(v)) => Some(*v),
546 _ => None,
547 })
548 };
549 let epoch = get("epoch").context("a timestamp struct with no epoch")?;
550 let nanos = match get("fraction") {
551 Some(fraction) => epoch
553 .checked_mul(1_000_000_000)
554 .and_then(|e| e.checked_add(fraction))
555 .context("timestamp out of range")?,
556 None => scaled(epoch, column.scale)?,
558 };
559 let offset = match column.kind {
560 Kind::TimestampTz => Some(zone(
561 get("timezone").context("a TIMESTAMP_TZ struct with no timezone")?,
562 )?),
563 _ => None,
564 };
565 Value::String(timestamp(nanos, column.scale, column.kind, offset)?)
566 }
567 (kind, _, physical) => bail!("a {kind:?} column sent as Arrow {physical:?}"),
568 })
569 })();
570 match decoded {
571 Ok(v) => v,
572 Err(e) => {
573 notes.warn(column, format!("{e:#}"));
574 raw_arrow(cell)
575 }
576 }
577}
578
579fn keys(columns: &[Column]) -> Vec<String> {
581 let mut seen: Vec<String> = Vec::with_capacity(columns.len());
582 for c in columns {
583 let mut key = c.name.clone();
584 let mut n = 1;
585 while seen.contains(&key) {
586 n += 1;
587 key = format!("{}_{n}", c.name);
588 }
589 seen.push(key);
590 }
591 seen
592}
593
594fn is_dml(statement_type_id: Option<i64>) -> bool {
596 statement_type_id.is_some_and(|id| (0x3000..0x4000).contains(&id))
597}
598
599pub fn decode(data: &Value, max_rows: usize) -> Result<Value> {
601 let rowtype = match data.get("rowtype") {
602 None | Some(Value::Null) => Vec::new(),
603 Some(Value::Array(a)) => a.clone(),
604 Some(_) => bail!("rowtype is not an array"),
605 };
606 ensure!(
607 rowtype.len() <= MAX_COLUMNS,
608 "the result has {} columns; at most {MAX_COLUMNS} are read",
609 rowtype.len()
610 );
611 let columns: Vec<Column> = rowtype
612 .iter()
613 .map(Column::from_rowtype)
614 .collect::<Result<_>>()?;
615 let names = keys(&columns);
616 let mut notes = Notes::default();
617 let format = data
618 .get("queryResultFormat")
619 .and_then(Value::as_str)
620 .unwrap_or("json")
621 .to_ascii_lowercase();
622
623 let mut rows: Vec<Value> = Vec::new();
624 let available: u64;
625 match format.as_str() {
626 "json" => {
627 let rowset = match data.get("rowset") {
628 None | Some(Value::Null) => Vec::new(),
629 Some(Value::Array(a)) => a.clone(),
630 Some(_) => bail!("rowset is not an array"),
631 };
632 available = rowset.len() as u64;
633 for raw in rowset.iter().take(max_rows) {
634 let cells = raw.as_array().context("a rowset row is not an array")?;
635 ensure!(
636 cells.len() == columns.len(),
637 "a row has {} cells where rowtype has {} columns",
638 cells.len(),
639 columns.len()
640 );
641 let mut row = Map::new();
642 for ((cell, column), key) in cells.iter().zip(&columns).zip(&names) {
643 row.insert(key.clone(), json_cell(cell, column, &mut notes));
644 }
645 rows.push(Value::Object(row));
646 }
647 }
648 "arrow" => {
649 let encoded = data
650 .get("rowsetBase64")
651 .and_then(Value::as_str)
652 .unwrap_or("");
653 if encoded.is_empty() {
654 available = 0;
655 } else {
656 use base64::Engine as _;
657 let stream = base64::engine::general_purpose::STANDARD
658 .decode(encoded.trim())
659 .context("rowsetBase64 is not base64")?;
660 let decoded = arrow::decode_stream(&stream, max_rows)?;
661 ensure!(
662 decoded.fields.len() == columns.len(),
663 "the Arrow schema has {} columns where rowtype has {}",
664 decoded.fields.len(),
665 columns.len()
666 );
667 available = decoded.total_rows;
668 for cells in &decoded.rows {
669 let mut row = Map::new();
670 for (((cell, field), column), key) in
671 cells.iter().zip(&decoded.fields).zip(&columns).zip(&names)
672 {
673 row.insert(key.clone(), arrow_cell(cell, field, column, &mut notes));
674 }
675 rows.push(Value::Object(row));
676 }
677 }
678 }
679 other => bail!("result format {other:?} is neither json nor arrow"),
680 }
681
682 let chunks = data
686 .get("chunks")
687 .and_then(Value::as_array)
688 .map(|c| c.len())
689 .unwrap_or(0);
690 let chunk_rows: u64 = data
691 .get("chunks")
692 .and_then(Value::as_array)
693 .map(|c| {
694 c.iter()
695 .filter_map(|c| c.get("rowCount").and_then(Value::as_u64))
696 .sum()
697 })
698 .unwrap_or(0);
699 let total = data.get("total").and_then(Value::as_u64);
700 let statement_type_id = data.get("statementTypeId").and_then(Value::as_i64);
701
702 let mut rows_affected: Option<i128> = None;
704 if is_dml(statement_type_id) || columns.iter().any(|c| c.name.starts_with("number of rows")) {
705 if let Some(Value::Object(first)) = rows.first() {
706 for (column, key) in columns.iter().zip(&names) {
707 if column.name.starts_with("number of rows")
708 || column.name.starts_with("number of multi-joined rows")
709 {
710 if let Some(n) = first.get(key).and_then(Value::as_i64) {
711 rows_affected = Some(rows_affected.unwrap_or(0) + n as i128);
712 }
713 }
714 }
715 }
716 }
717
718 let delivered = rows.len() as u64;
719 let truncated = delivered < available || chunks > 0;
720 Ok(json!({
721 "query_id": data.get("queryId").and_then(Value::as_str).unwrap_or(""),
722 "statement_type_id": statement_type_id,
723 "result_format": format,
724 "columns": columns.iter().map(Column::to_json).collect::<Vec<_>>(),
725 "row_count": delivered,
726 "rows": rows,
727 "total": total,
728 "rows_affected": rows_affected.map(|n| n as i64),
729 "truncated": truncated,
730 "chunks_not_fetched": chunks,
731 "rows_in_chunks": chunk_rows,
732 "decode_warnings": notes.warnings,
733 "decode_warning_count": notes.warning_count,
734 "truncated_cells": notes.truncated_cells,
735 }))
736}