Skip to main content

scopedb_client/
result.rs

1// Copyright 2024 ScopeDB, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::collections::HashSet;
17use std::fmt;
18use std::str::FromStr;
19
20use crate::DataType;
21use crate::Error;
22use crate::ErrorKind;
23use crate::protocol::ResultSetData;
24use crate::protocol::StatementResultSet;
25
26#[derive(Debug, Clone)]
27pub struct Schema {
28    pub(crate) fields: Vec<FieldSchema>,
29}
30
31impl Schema {
32    pub fn fields(&self) -> &[FieldSchema] {
33        &self.fields
34    }
35}
36
37#[derive(Debug, Clone)]
38pub struct FieldSchema {
39    pub(crate) name: String,
40    pub(crate) data_type: DataType,
41}
42
43impl FieldSchema {
44    pub fn name(&self) -> &str {
45        &self.name
46    }
47
48    pub fn data_type(&self) -> DataType {
49        self.data_type
50    }
51}
52
53#[derive(Debug, Clone)]
54pub struct ResultSet {
55    schema: Schema,
56    num_rows: usize,
57    data: ResultSetData,
58}
59
60/// One typed result row keyed by its ScopeQL output column name.
61pub type ResultObject = BTreeMap<String, Value>;
62
63impl ResultSet {
64    pub fn num_rows(&self) -> usize {
65        self.num_rows
66    }
67
68    pub fn schema(&self) -> &Schema {
69        &self.schema
70    }
71
72    /// Returns the unparsed JSON-format cells exactly as received from ScopeDB.
73    pub fn raw_rows(&self) -> &[Vec<Option<String>>] {
74        match &self.data {
75            ResultSetData::Json { rows } => rows,
76        }
77    }
78
79    /// Converts all result cells to their typed Rust representation.
80    pub fn to_values(&self) -> Result<Vec<Vec<Value>>, Error> {
81        self.validate_shape()?;
82        self.raw_rows()
83            .iter()
84            .enumerate()
85            .map(|(index, row)| {
86                row.iter()
87                    .zip(self.schema.fields.iter())
88                    .map(|(cell, field)| {
89                        parse_cell(cell.as_deref(), field.data_type()).map_err(|error| {
90                            error
91                                .with_context("row", index)
92                                .with_context("column", field.name())
93                        })
94                    })
95                    .collect()
96            })
97            .collect()
98    }
99
100    /// Consumes the result set and returns typed rows.
101    pub fn into_values(self) -> Result<Vec<Vec<Value>>, Error> {
102        self.validate_shape()?;
103        let ResultSet { schema, data, .. } = self;
104        let rows = match data {
105            ResultSetData::Json { rows } => rows,
106        };
107        rows.into_iter()
108            .enumerate()
109            .map(|(index, row)| {
110                row.into_iter()
111                    .zip(schema.fields.iter())
112                    .map(|(cell, field)| {
113                        parse_owned_cell(cell, field.data_type()).map_err(|error| {
114                            error
115                                .with_context("row", index)
116                                .with_context("column", field.name())
117                        })
118                    })
119                    .collect()
120            })
121            .collect()
122    }
123
124    /// Converts all rows to maps keyed by unique output column names.
125    ///
126    /// Duplicate output names are rejected instead of silently discarding a
127    /// value. Use [`ResultSet::to_values`] when a query intentionally returns
128    /// duplicate column names.
129    pub fn to_objects(&self) -> Result<Vec<ResultObject>, Error> {
130        self.validate_unique_fields()?;
131        self.to_values().map(|rows| {
132            rows.into_iter()
133                .map(|row| {
134                    self.schema
135                        .fields
136                        .iter()
137                        .map(|field| field.name.clone())
138                        .zip(row)
139                        .collect()
140                })
141                .collect()
142        })
143    }
144
145    /// Consumes the result set and returns rows keyed by output column name.
146    pub fn into_objects(self) -> Result<Vec<ResultObject>, Error> {
147        self.validate_unique_fields()?;
148        let names = self
149            .schema
150            .fields
151            .iter()
152            .map(|field| field.name.clone())
153            .collect::<Vec<_>>();
154        self.into_values().map(|rows| {
155            rows.into_iter()
156                .map(|values| names.iter().cloned().zip(values).collect())
157                .collect()
158        })
159    }
160
161    /// Returns the first row keyed by output column name, or `None` when empty.
162    pub fn first(&self) -> Result<Option<ResultObject>, Error> {
163        self.validate_row_count()?;
164        let Some(row) = self.raw_rows().first() else {
165            return Ok(None);
166        };
167        self.validate_unique_fields()?;
168        self.validate_row_shape(0, row)?;
169        let values = row
170            .iter()
171            .zip(self.schema.fields.iter())
172            .map(|(cell, field)| {
173                parse_cell(cell.as_deref(), field.data_type()).map_err(|error| {
174                    error
175                        .with_context("row", 0)
176                        .with_context("column", field.name())
177                })
178            })
179            .collect::<Result<Vec<_>, _>>()?;
180        Ok(Some(
181            self.schema
182                .fields
183                .iter()
184                .map(|field| field.name.clone())
185                .zip(values)
186                .collect(),
187        ))
188    }
189
190    fn validate_row_count(&self) -> Result<(), Error> {
191        let rows = self.raw_rows();
192        if rows.len() != self.num_rows {
193            return Err(Error::new(
194                ErrorKind::Unexpected,
195                format!(
196                    "result row count mismatch: expected {}, got {}",
197                    self.num_rows,
198                    rows.len()
199                ),
200            ));
201        }
202        Ok(())
203    }
204
205    fn validate_shape(&self) -> Result<(), Error> {
206        self.validate_row_count()?;
207        let rows = self.raw_rows();
208        let expected = self.schema.fields.len();
209        if let Some((index, row)) = rows
210            .iter()
211            .enumerate()
212            .find(|(_, row)| row.len() != expected)
213        {
214            return self.validate_row_shape(index, row);
215        }
216        Ok(())
217    }
218
219    fn validate_row_shape(&self, index: usize, row: &[Option<String>]) -> Result<(), Error> {
220        let expected = self.schema.fields.len();
221        if row.len() == expected {
222            return Ok(());
223        }
224        Err(Error::new(
225            ErrorKind::Unexpected,
226            format!(
227                "result row {index} field count mismatch: expected {expected}, got {}",
228                row.len()
229            ),
230        ))
231    }
232
233    fn validate_unique_fields(&self) -> Result<(), Error> {
234        let mut names = HashSet::with_capacity(self.schema.fields.len());
235        if let Some(field) = self
236            .schema
237            .fields
238            .iter()
239            .find(|field| !names.insert(field.name.as_str()))
240        {
241            return Err(Error::new(
242                ErrorKind::Unexpected,
243                format!(
244                    "cannot convert result to objects because column name {:?} is duplicated; use to_values instead",
245                    field.name
246                ),
247            ));
248        }
249        Ok(())
250    }
251
252    pub(crate) fn from_statement_result_set(result_set: StatementResultSet) -> ResultSet {
253        ResultSet {
254            num_rows: result_set.metadata.num_rows,
255            schema: Schema {
256                fields: result_set
257                    .metadata
258                    .fields
259                    .into_iter()
260                    .map(|field| FieldSchema {
261                        name: field.name,
262                        data_type: field.data_type,
263                    })
264                    .collect(),
265            },
266            data: result_set.data,
267        }
268    }
269}
270
271fn parse_cell(value: Option<&str>, data_type: DataType) -> Result<Value, Error> {
272    let Some(value) = value else {
273        return Ok(Value::Null);
274    };
275
276    match data_type {
277        DataType::Int => i64::from_str(value).map(Value::Int).map_err(|err| {
278            Error::new(
279                ErrorKind::Unexpected,
280                format!("failed to parse int value: {err}"),
281            )
282        }),
283        DataType::UInt => u64::from_str(value).map(Value::UInt).map_err(|err| {
284            Error::new(
285                ErrorKind::Unexpected,
286                format!("failed to parse uint value: {err}"),
287            )
288        }),
289        DataType::Float => {
290            let parsed = f64::from_str(value).map_err(|err| {
291                Error::new(
292                    ErrorKind::Unexpected,
293                    format!("failed to parse float value: {err}"),
294                )
295            })?;
296            if !parsed.is_finite() && !matches!(value, "NaN" | "inf" | "-inf") {
297                return Err(Error::new(
298                    ErrorKind::Unexpected,
299                    "float value is outside the f64 range",
300                ));
301            }
302            Ok(Value::Float(parsed))
303        }
304        DataType::Timestamp => jiff::Timestamp::from_str(value)
305            .map(Value::Timestamp)
306            .map_err(|err| {
307                Error::new(
308                    ErrorKind::Unexpected,
309                    format!("failed to parse timestamp value: {err}"),
310                )
311            }),
312        DataType::Interval => jiff::SignedDuration::from_str(value)
313            .map(Value::Interval)
314            .map_err(|err| {
315                Error::new(
316                    ErrorKind::Unexpected,
317                    format!("failed to parse interval value: {err}"),
318                )
319            }),
320        DataType::Boolean => bool::from_str(value).map(Value::Boolean).map_err(|err| {
321            Error::new(
322                ErrorKind::Unexpected,
323                format!("failed to parse boolean value: {err}"),
324            )
325        }),
326        DataType::String => Ok(Value::String(value.to_string())),
327        DataType::Binary => Ok(Value::Binary(value.to_string())),
328        DataType::Array => Ok(Value::Array(value.to_string())),
329        DataType::Object => Ok(Value::Object(value.to_string())),
330        DataType::Any => Ok(Value::Any(value.to_string())),
331        DataType::Null => Err(Error::new(
332            ErrorKind::Unexpected,
333            "non-null cell returned for a null column",
334        )),
335    }
336}
337
338fn parse_owned_cell(value: Option<String>, data_type: DataType) -> Result<Value, Error> {
339    let Some(value) = value else {
340        return Ok(Value::Null);
341    };
342
343    match data_type {
344        DataType::String => Ok(Value::String(value)),
345        DataType::Binary => Ok(Value::Binary(value)),
346        DataType::Array => Ok(Value::Array(value)),
347        DataType::Object => Ok(Value::Object(value)),
348        DataType::Any => Ok(Value::Any(value)),
349        _ => parse_cell(Some(&value), data_type),
350    }
351}
352
353/// A typed result cell. Equality compares encoded array, object, any, and binary
354/// values as strings; floats use Rust's usual equality, including NaN != NaN.
355#[derive(Clone, PartialEq)]
356pub enum Value {
357    /// Signed integer value.
358    Int(i64),
359    /// Unsigned integer value.
360    UInt(u64),
361    /// Float value.
362    Float(f64),
363    /// Timestamp value.
364    Timestamp(jiff::Timestamp),
365    /// Interval value.
366    Interval(jiff::SignedDuration),
367    /// Boolean value.
368    Boolean(bool),
369    /// String value.
370    String(String),
371    /// Binary data represented as a hex string.
372    Binary(String),
373    /// Array of values, represented in its string format.
374    Array(String),
375    /// Object represented in its string format.
376    Object(String),
377    /// Any value, represented in its string format.
378    Any(String),
379    /// Null value.
380    Null,
381}
382
383impl fmt::Debug for Value {
384    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
385        match self {
386            Value::Int(v) => write!(f, "{v}"),
387            Value::UInt(v) => write!(f, "{v}"),
388            Value::Float(v) => write!(f, "{v:?}"),
389            Value::Timestamp(v) => format_timestamp(f, v),
390            Value::Interval(v) => format_interval(f, v),
391            Value::Boolean(v) => write!(f, "{v}"),
392            Value::String(v) => quote_string(f, v, '\''),
393            Value::Binary(v) => write!(f, "{v}"),
394            Value::Array(v) => write!(f, "{v}"),
395            Value::Object(v) => write!(f, "{v}"),
396            Value::Any(v) => write!(f, "{v}"),
397            Value::Null => write!(f, "NULL"),
398        }
399    }
400}
401
402impl fmt::Display for Value {
403    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
404        match self {
405            Value::Int(v) => write!(f, "{v}"),
406            Value::UInt(v) => write!(f, "{v}"),
407            Value::Float(v) => write!(f, "{v:?}"),
408            Value::Timestamp(v) => format_timestamp(f, v),
409            Value::Interval(v) => format_interval(f, v),
410            Value::Boolean(v) => write!(f, "{v}"),
411            Value::String(v) => write!(f, "{v}"),
412            Value::Binary(v) => write!(f, "{v}"),
413            Value::Array(v) => write!(f, "{v}"),
414            Value::Object(v) => write!(f, "{v}"),
415            Value::Any(v) => write!(f, "{v}"),
416            Value::Null => write!(f, "NULL"),
417        }
418    }
419}
420
421fn format_timestamp(f: &mut fmt::Formatter<'_>, ts: &jiff::Timestamp) -> fmt::Result {
422    use jiff::fmt::StdFmtWrite;
423    use jiff::fmt::temporal;
424
425    let precision = f.precision().map(|p| u8::try_from(p).unwrap_or(u8::MAX));
426    temporal::DateTimePrinter::new()
427        .precision(precision)
428        .print_timestamp(ts, StdFmtWrite(f))
429        .map_err(|_| fmt::Error)
430}
431
432fn format_interval(f: &mut fmt::Formatter<'_>, d: &jiff::SignedDuration) -> fmt::Result {
433    use jiff::fmt::StdFmtWrite;
434    use jiff::fmt::friendly;
435
436    friendly::SpanPrinter::new()
437        .spacing(friendly::Spacing::None)
438        .designator(friendly::Designator::Compact)
439        .print_duration(d, StdFmtWrite(f))
440        .map_err(|_| fmt::Error)
441}
442
443fn quote_string(f: &mut fmt::Formatter<'_>, s: &str, quote: char) -> fmt::Result {
444    write!(f, "{quote}")?;
445    for c in s.chars() {
446        match c {
447            '\t' => write!(f, "\\t")?,
448            '\r' => write!(f, "\\r")?,
449            '\n' => write!(f, "\\n")?,
450            '\\' => write!(f, "\\\\")?,
451            '\x00'..='\x1F' => write!(f, "\\x{:02x}", c as u8)?,
452            c => {
453                if c != quote {
454                    write!(f, "{c}")?
455                } else {
456                    write!(f, "\\{quote}")?
457                }
458            }
459        }
460    }
461    write!(f, "{quote}")
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use crate::protocol::FieldMetadata;
468    use crate::protocol::ResultSetMetadata;
469    use crate::protocol::StatementResultSet;
470
471    fn result(fields: Vec<(&str, DataType)>, rows: Vec<Vec<Option<&str>>>) -> ResultSet {
472        let num_rows = rows.len();
473        ResultSet::from_statement_result_set(StatementResultSet {
474            metadata: ResultSetMetadata {
475                fields: fields
476                    .into_iter()
477                    .map(|(name, data_type)| FieldMetadata {
478                        name: name.to_string(),
479                        data_type,
480                    })
481                    .collect(),
482                num_rows,
483            },
484            data: ResultSetData::Json {
485                rows: rows
486                    .into_iter()
487                    .map(|row| {
488                        row.into_iter()
489                            .map(|cell| cell.map(str::to_string))
490                            .collect()
491                    })
492                    .collect(),
493            },
494        })
495    }
496
497    #[test]
498    fn exposes_raw_values_objects_and_first_row() {
499        let result = result(
500            vec![("id", DataType::Int), ("name", DataType::String)],
501            vec![vec![Some("42"), Some("ScopeDB")]],
502        );
503
504        assert_eq!(result.raw_rows()[0][0].as_deref(), Some("42"));
505        let values = result.to_values().unwrap();
506        assert!(matches!(values[0][0], Value::Int(42)));
507        let objects = result.to_objects().unwrap();
508        assert!(matches!(objects[0].get("id"), Some(Value::Int(42))));
509        assert!(matches!(
510            result.first().unwrap().unwrap().get("name"),
511            Some(Value::String(value)) if value == "ScopeDB"
512        ));
513    }
514
515    #[test]
516    fn consuming_conversions_move_string_values() {
517        let values = result(
518            vec![("name", DataType::String)],
519            vec![vec![Some("ScopeDB")]],
520        )
521        .into_values()
522        .unwrap();
523        assert!(matches!(&values[0][0], Value::String(value) if value == "ScopeDB"));
524
525        let objects = result(
526            vec![("name", DataType::String)],
527            vec![vec![Some("ScopeDB")]],
528        )
529        .into_objects()
530        .unwrap();
531        assert!(matches!(
532            objects[0].get("name"),
533            Some(Value::String(value)) if value == "ScopeDB"
534        ));
535    }
536
537    #[test]
538    fn malformed_shapes_return_errors_instead_of_panicking() {
539        let mut malformed = result(
540            vec![("id", DataType::Int)],
541            vec![vec![Some("1"), Some("extra")]],
542        );
543        malformed.num_rows = 2;
544        let error = malformed.to_values().unwrap_err();
545        assert_eq!(error.kind(), ErrorKind::Unexpected);
546    }
547
548    #[test]
549    fn object_conversion_rejects_duplicate_names() {
550        let result = result(
551            vec![("value", DataType::Int), ("value", DataType::Int)],
552            vec![vec![Some("1"), Some("2")]],
553        );
554        let error = result.to_objects().unwrap_err();
555        assert!(error.message().contains("duplicated"));
556    }
557
558    #[test]
559    fn first_parses_only_the_first_row() {
560        let result = result(
561            vec![("id", DataType::Int)],
562            vec![vec![Some("1")], vec![Some("2"), Some("extra")]],
563        );
564        assert!(matches!(
565            result.first().unwrap().unwrap().get("id"),
566            Some(Value::Int(1))
567        ));
568        assert!(result.to_values().is_err());
569    }
570
571    #[test]
572    fn first_returns_none_for_empty_results_with_duplicate_names() {
573        let result = result(
574            vec![("value", DataType::Int), ("value", DataType::Int)],
575            vec![],
576        );
577        assert!(result.first().unwrap().is_none());
578    }
579
580    #[test]
581    fn numeric_boundaries_and_raw_json_are_lossless() {
582        for (kind, wire) in [
583            (DataType::Int, "-9223372036854775808"),
584            (DataType::Int, "9223372036854775807"),
585            (DataType::Int, "9007199254740993"),
586            (DataType::UInt, "18446744073709551615"),
587            (DataType::Object, r#"{"id":184467440737095516160}"#),
588            (DataType::Array, "[0.12345678901234567890123456789]"),
589        ] {
590            let result = result(vec![("value", kind)], vec![vec![Some(wire)]]);
591            assert_eq!(result.raw_rows()[0][0].as_deref(), Some(wire));
592            assert_eq!(result.to_values().unwrap()[0][0].to_string(), wire);
593            assert_eq!(
594                result.clone().into_objects().unwrap()[0]["value"].to_string(),
595                wire
596            );
597            assert_eq!(result.first().unwrap().unwrap()["value"].to_string(), wire);
598        }
599        for (kind, wire) in [
600            (DataType::Int, "9223372036854775808"),
601            (DataType::Int, "-9223372036854775809"),
602            (DataType::UInt, "18446744073709551616"),
603            (DataType::UInt, "-1"),
604            (DataType::Float, "1e309"),
605        ] {
606            let result = result(vec![("value", kind)], vec![vec![Some(wire)]]);
607            for error in [
608                result.to_values().unwrap_err(),
609                result.clone().into_values().unwrap_err(),
610                result.clone().into_objects().unwrap_err(),
611                result.first().unwrap_err(),
612            ] {
613                let diagnostic = error.to_string();
614                assert!(diagnostic.contains("row: 0") && diagnostic.contains("column: value"));
615            }
616        }
617    }
618
619    #[test]
620    fn float_boundaries_preserve_server_special_values() {
621        for (wire, expected) in [
622            ("1.7976931348623157e308", f64::MAX),
623            ("5e-324", f64::from_bits(1)),
624            ("inf", f64::INFINITY),
625            ("-inf", f64::NEG_INFINITY),
626        ] {
627            let Value::Float(value) = parse_cell(Some(wire), DataType::Float).unwrap() else {
628                panic!("expected float")
629            };
630            assert_eq!(value, expected);
631        }
632        assert!(
633            matches!(parse_cell(Some("NaN"), DataType::Float).unwrap(), Value::Float(value) if value.is_nan())
634        );
635    }
636
637    #[test]
638    fn first_rejects_missing_rows_advertised_by_metadata() {
639        let mut result = result(vec![("id", DataType::Int)], vec![]);
640        result.num_rows = 1;
641        assert!(
642            result
643                .first()
644                .unwrap_err()
645                .message()
646                .contains("row count mismatch")
647        );
648    }
649
650    #[test]
651    fn binary_display_preserves_the_wire_hex_string() {
652        let value = Value::Binary("00FF".to_string());
653        assert_eq!(value.to_string(), "00FF");
654        assert_eq!(format!("{value:?}"), "00FF");
655    }
656}