Skip to main content

omni_dev/snowflake/client/
row.rs

1//! Result rows and arbitrary-schema → JSON decoding.
2//!
3//! Snowflake's v1 query endpoint returns column metadata (`rowtype`) plus a
4//! row set whose cells are stringified (or null). This module models a column,
5//! a row, and the mapping from a raw cell to a `serde_json::Value` driven by the
6//! column's Snowflake type tag — implemented from the documented wire formats
7//! (e.g. `DATE` is days-since-epoch, timestamps are `seconds[.fraction]`).
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use chrono::{DateTime, Days, NaiveDate, NaiveTime, Utc};
13use serde_json::{json, Map, Value};
14
15/// Metadata for one result column, from the query response `rowtype`.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct Column {
18    /// Column name as Snowflake reported it.
19    pub name: String,
20    /// Lowercase Snowflake type tag (e.g. `fixed`, `text`, `timestamp_ntz`).
21    pub ty: String,
22    /// Whether the column may be null.
23    pub nullable: bool,
24    /// Declared length (text), when reported.
25    pub length: Option<i64>,
26    /// Declared precision (number), when reported.
27    pub precision: Option<i64>,
28    /// Declared scale (number/time/timestamp), when reported.
29    pub scale: Option<i64>,
30}
31
32impl Column {
33    /// A compact, lowercase type label with parameters when reported (e.g.
34    /// `fixed(10,2)`, `text(255)`, `timestamp_ntz(9)`).
35    #[must_use]
36    pub fn type_label(&self) -> String {
37        match (self.precision, self.scale, self.length) {
38            (Some(p), Some(s), _) => format!("{}({p},{s})", self.ty),
39            (Some(p), None, _) => format!("{}({p})", self.ty),
40            (None, Some(s), _) => format!("{}({s})", self.ty),
41            (None, None, Some(l)) => format!("{}({l})", self.ty),
42            _ => self.ty.clone(),
43        }
44    }
45}
46
47/// One result row: raw stringified cells plus the shared column schema.
48#[derive(Clone, Debug)]
49pub struct Row {
50    values: Vec<Option<String>>,
51    columns: Arc<Vec<Column>>,
52    index: Arc<HashMap<String, usize>>,
53}
54
55impl Row {
56    /// Builds a row from its raw cells and the shared schema/index.
57    #[must_use]
58    pub fn new(
59        values: Vec<Option<String>>,
60        columns: Arc<Vec<Column>>,
61        index: Arc<HashMap<String, usize>>,
62    ) -> Self {
63        Self {
64            values,
65            columns,
66            index,
67        }
68    }
69
70    /// The shared column schema.
71    #[must_use]
72    pub fn columns(&self) -> &[Column] {
73        &self.columns
74    }
75
76    /// The raw cell value at `index`, if present.
77    #[must_use]
78    pub fn raw_at(&self, index: usize) -> Option<&str> {
79        self.values.get(index).and_then(|v| v.as_deref())
80    }
81
82    /// The raw cell value for a column name (case-insensitive), if present.
83    #[must_use]
84    pub fn raw(&self, name: &str) -> Option<&str> {
85        self.index
86            .get(&name.to_ascii_uppercase())
87            .and_then(|&i| self.raw_at(i))
88    }
89
90    /// Converts the row to a JSON object keyed by column name.
91    ///
92    /// Repeated column names (e.g. `SELECT 1, 1`) are disambiguated with a
93    /// `_<n>` suffix so no column is silently dropped — JSON object keys must be
94    /// unique.
95    #[must_use]
96    pub fn to_json_object(&self) -> Map<String, Value> {
97        let mut map = Map::with_capacity(self.columns.len());
98        let mut seen: HashMap<&str, u32> = HashMap::new();
99        for (i, col) in self.columns.iter().enumerate() {
100            let raw = self.values.get(i).and_then(|v| v.as_deref());
101            let value = value_to_json(raw, col);
102            let count = seen.entry(col.name.as_str()).or_insert(0);
103            *count += 1;
104            let key = if *count == 1 {
105                col.name.clone()
106            } else {
107                format!("{}_{}", col.name, *count)
108            };
109            map.insert(key, value);
110        }
111        map
112    }
113}
114
115/// Builds the self-describing `{ columns: [{name, type}], rows: [{col: val}] }`
116/// payload from a query's rows. Column metadata comes from the first row (an
117/// empty result reports `columns: []`).
118#[must_use]
119pub fn rows_to_payload(rows: &[Row]) -> Value {
120    let columns: Vec<Value> = rows.first().map_or_else(Vec::new, |row| {
121        row.columns()
122            .iter()
123            .map(|c| json!({ "name": c.name, "type": c.type_label() }))
124            .collect()
125    });
126    let out: Vec<Value> = rows
127        .iter()
128        .map(|r| Value::Object(r.to_json_object()))
129        .collect();
130    json!({ "columns": columns, "rows": out })
131}
132
133/// Maps a single raw cell to JSON according to its column type. Total: any
134/// decode failure falls back to the raw string (or `null`).
135#[must_use]
136pub fn value_to_json(raw: Option<&str>, col: &Column) -> Value {
137    let Some(s) = raw else { return Value::Null };
138    match col.ty.as_str() {
139        "fixed" if col.scale.unwrap_or(0) == 0 => s
140            .parse::<i64>()
141            .map_or_else(|_| Value::String(s.to_string()), |n| json!(n)),
142        "fixed" | "real" | "float" | "double" | "double precision" => number_or_string(s),
143        "boolean" => match s {
144            "1" => json!(true),
145            "0" => json!(false),
146            _ if s.eq_ignore_ascii_case("true") => json!(true),
147            _ if s.eq_ignore_ascii_case("false") => json!(false),
148            _ => Value::String(s.to_string()),
149        },
150        "date" => decode_date(s),
151        "time" => decode_time(s),
152        "timestamp_ntz" => decode_epoch(s).map_or_else(
153            || Value::String(s.to_string()),
154            |dt| json!(dt.naive_utc().format("%Y-%m-%dT%H:%M:%S%.f").to_string()),
155        ),
156        "timestamp_ltz" => decode_epoch(s)
157            .map_or_else(|| Value::String(s.to_string()), |dt| json!(dt.to_rfc3339())),
158        "timestamp_tz" => decode_timestamp_tz(s),
159        "variant" | "object" | "array" => {
160            serde_json::from_str(s).unwrap_or_else(|_| Value::String(s.to_string()))
161        }
162        // text/varchar/char/string, binary (hex), geography, geometry, vector,
163        // and any unknown type: the raw wire string.
164        _ => Value::String(s.to_string()),
165    }
166}
167
168/// Parses a float string to a JSON number, falling back to the raw string for
169/// values JSON cannot represent.
170fn number_or_string(s: &str) -> Value {
171    s.parse::<f64>()
172        .ok()
173        .and_then(serde_json::Number::from_f64)
174        .map_or_else(|| Value::String(s.to_string()), Value::Number)
175}
176
177/// `DATE` is days since the Unix epoch.
178fn decode_date(s: &str) -> Value {
179    let Ok(days) = s.parse::<i64>() else {
180        return Value::String(s.to_string());
181    };
182    let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap_or_default();
183    let date = if days >= 0 {
184        epoch.checked_add_days(Days::new(days.unsigned_abs()))
185    } else {
186        epoch.checked_sub_days(Days::new(days.unsigned_abs()))
187    };
188    date.map_or_else(
189        || Value::String(s.to_string()),
190        |d| json!(d.format("%Y-%m-%d").to_string()),
191    )
192}
193
194/// `TIME` is seconds-from-midnight with an optional fractional part.
195fn decode_time(s: &str) -> Value {
196    let Some((secs, nanos)) = split_seconds(s) else {
197        return Value::String(s.to_string());
198    };
199    if secs < 0 {
200        return Value::String(s.to_string());
201    }
202    NaiveTime::from_num_seconds_from_midnight_opt(secs.unsigned_abs() as u32, nanos).map_or_else(
203        || Value::String(s.to_string()),
204        |t| json!(t.format("%H:%M:%S%.f").to_string()),
205    )
206}
207
208/// `TIMESTAMP_TZ` wire form is `"<epoch> <offset-minutes-plus-1440>"`; we keep
209/// the UTC instant (the offset is advisory and dropped here).
210fn decode_timestamp_tz(s: &str) -> Value {
211    let epoch = s.split_whitespace().next().unwrap_or(s);
212    decode_epoch(epoch).map_or_else(|| Value::String(s.to_string()), |dt| json!(dt.to_rfc3339()))
213}
214
215/// Parses an `seconds[.fraction]` epoch value into a UTC instant.
216fn decode_epoch(s: &str) -> Option<DateTime<Utc>> {
217    let (secs, nanos) = split_seconds(s)?;
218    DateTime::from_timestamp(secs, nanos)
219}
220
221/// Splits a signed `seconds[.fraction]` string into `(seconds, nanoseconds)`,
222/// flooring toward negative infinity so fractional pre-epoch values are correct.
223fn split_seconds(s: &str) -> Option<(i64, u32)> {
224    let s = s.trim();
225    let negative = s.starts_with('-');
226    let (secs_str, frac_str) = s.split_once('.').unwrap_or((s, ""));
227    let mut secs = secs_str.parse::<i64>().ok()?;
228    if frac_str.is_empty() {
229        return Some((secs, 0));
230    }
231    if !frac_str.bytes().all(|b| b.is_ascii_digit()) {
232        return None;
233    }
234    // Pad/truncate the fraction to exactly 9 digits (nanoseconds).
235    let mut digits: Vec<u8> = frac_str.bytes().take(9).collect();
236    digits.resize(9, b'0');
237    let mut nanos = std::str::from_utf8(&digits).ok()?.parse::<u32>().ok()?;
238    if negative && nanos > 0 {
239        // "-1.5s" = floor(-1.5)=-2 plus 0.5s; "-0.x" parses secs as 0.
240        secs -= 1;
241        nanos = 1_000_000_000 - nanos;
242    }
243    Some((secs, nanos))
244}
245
246#[cfg(test)]
247#[allow(clippy::unwrap_used, clippy::expect_used)]
248mod tests {
249    use super::*;
250
251    fn col(ty: &str, precision: Option<i64>, scale: Option<i64>, length: Option<i64>) -> Column {
252        Column {
253            name: "C".to_string(),
254            ty: ty.to_string(),
255            nullable: true,
256            length,
257            precision,
258            scale,
259        }
260    }
261
262    #[test]
263    fn type_label_renders_parameters() {
264        assert_eq!(
265            col("fixed", Some(10), Some(2), None).type_label(),
266            "fixed(10,2)"
267        );
268        assert_eq!(col("text", None, None, Some(255)).type_label(), "text(255)");
269        assert_eq!(
270            col("timestamp_ntz", None, None, None).type_label(),
271            "timestamp_ntz"
272        );
273    }
274
275    #[test]
276    fn null_is_json_null() {
277        assert_eq!(
278            value_to_json(None, &col("fixed", None, Some(0), None)),
279            Value::Null
280        );
281    }
282
283    #[test]
284    fn decodes_each_type() {
285        assert_eq!(
286            value_to_json(Some("42"), &col("fixed", Some(38), Some(0), None)),
287            json!(42)
288        );
289        let big = "123456789012345678901234567890";
290        assert_eq!(
291            value_to_json(Some(big), &col("fixed", Some(38), Some(0), None)),
292            json!(big),
293            "overflowing integer keeps the exact string"
294        );
295        assert_eq!(
296            value_to_json(Some("123.45"), &col("fixed", Some(10), Some(2), None)),
297            json!(123.45)
298        );
299        assert_eq!(
300            value_to_json(Some("1.5"), &col("real", None, None, None)),
301            json!(1.5)
302        );
303        assert_eq!(
304            value_to_json(Some("1"), &col("boolean", None, None, None)),
305            json!(true)
306        );
307        assert_eq!(
308            value_to_json(Some("0"), &col("boolean", None, None, None)),
309            json!(false)
310        );
311        assert_eq!(
312            value_to_json(Some("hello"), &col("text", None, None, Some(255))),
313            json!("hello")
314        );
315        assert_eq!(
316            value_to_json(Some("0"), &col("date", None, None, None)),
317            json!("1970-01-01")
318        );
319        assert_eq!(
320            value_to_json(Some("19723"), &col("date", None, None, None)),
321            json!("2024-01-01")
322        );
323        assert_eq!(
324            value_to_json(Some("3661"), &col("time", None, None, Some(0))),
325            json!("01:01:01")
326        );
327        assert_eq!(
328            value_to_json(
329                Some("1704067200"),
330                &col("timestamp_ntz", None, None, Some(0))
331            ),
332            json!("2024-01-01T00:00:00")
333        );
334        assert_eq!(
335            value_to_json(
336                Some("1704067200"),
337                &col("timestamp_ltz", None, None, Some(0))
338            ),
339            json!("2024-01-01T00:00:00+00:00")
340        );
341        assert_eq!(
342            value_to_json(
343                Some("1704067200.000 1440"),
344                &col("timestamp_tz", None, None, Some(3))
345            ),
346            json!("2024-01-01T00:00:00+00:00")
347        );
348        assert_eq!(
349            value_to_json(Some(r#"{"a":1}"#), &col("variant", None, None, None)),
350            json!({ "a": 1 })
351        );
352        assert_eq!(
353            value_to_json(Some("DEADBEEF"), &col("binary", None, None, None)),
354            json!("DEADBEEF")
355        );
356    }
357
358    #[test]
359    fn malformed_value_falls_back_to_raw() {
360        assert_eq!(
361            value_to_json(Some("maybe"), &col("boolean", None, None, None)),
362            json!("maybe")
363        );
364    }
365
366    #[test]
367    fn row_to_json_keys_by_column_name() {
368        let columns = Arc::new(vec![
369            Column {
370                name: "ID".to_string(),
371                ty: "fixed".to_string(),
372                nullable: false,
373                length: None,
374                precision: Some(38),
375                scale: Some(0),
376            },
377            Column {
378                name: "NAME".to_string(),
379                ty: "text".to_string(),
380                nullable: true,
381                length: None,
382                precision: None,
383                scale: None,
384            },
385        ]);
386        let mut index = HashMap::new();
387        index.insert("ID".to_string(), 0);
388        index.insert("NAME".to_string(), 1);
389        let row = Row::new(
390            vec![Some("1".to_string()), Some("hi".to_string())],
391            columns,
392            Arc::new(index),
393        );
394        let obj = row.to_json_object();
395        assert_eq!(obj.get("ID"), Some(&json!(1)));
396        assert_eq!(obj.get("NAME"), Some(&json!("hi")));
397        assert_eq!(obj.len(), 2);
398    }
399
400    #[test]
401    fn duplicate_column_names_are_disambiguated() {
402        let dup = || Column {
403            name: "N".to_string(),
404            ty: "fixed".to_string(),
405            nullable: false,
406            length: None,
407            precision: Some(38),
408            scale: Some(0),
409        };
410        let columns = Arc::new(vec![dup(), dup()]);
411        let mut index = HashMap::new();
412        index.insert("N".to_string(), 0);
413        let row = Row::new(
414            vec![Some("1".to_string()), Some("2".to_string())],
415            columns,
416            Arc::new(index),
417        );
418        let obj = row.to_json_object();
419        assert_eq!(obj.get("N"), Some(&json!(1)));
420        assert_eq!(obj.get("N_2"), Some(&json!(2)), "second N is not dropped");
421        assert_eq!(obj.len(), 2);
422    }
423}