gcloud_spanner/
row.rs

1use std::collections::{BTreeMap, HashMap};
2use std::num::ParseIntError;
3use std::str::FromStr;
4use std::sync::Arc;
5
6use base64::prelude::*;
7use base64::DecodeError;
8use prost_types::value::Kind;
9use prost_types::{value, Value};
10use time::format_description::well_known::Rfc3339;
11use time::macros::format_description;
12use time::{Date, OffsetDateTime};
13
14use google_cloud_googleapis::spanner::v1::struct_type::Field;
15use google_cloud_googleapis::spanner::v1::StructType;
16
17use crate::bigdecimal::{BigDecimal, ParseBigDecimalError};
18use crate::value::CommitTimestamp;
19
20#[derive(Clone)]
21pub struct Row {
22    index: Arc<HashMap<String, usize>>,
23    fields: Arc<Vec<Field>>,
24    values: Vec<Value>,
25}
26
27#[derive(thiserror::Error, Debug)]
28pub enum Error {
29    #[error("Illegal Kind: field={0}, kind={1}")]
30    KindMismatch(String, String),
31    #[error("No kind found: field={0}")]
32    NoKind(String),
33    #[error("Parse field: field={0}")]
34    IntParseError(String, #[source] ParseIntError),
35    #[error("Failed to parse as Date|DateTime {0}")]
36    DateParseError(String, #[source] time::error::Parse),
37    #[error("Failed to parse as ByteArray {0}")]
38    ByteParseError(String, #[source] DecodeError),
39    #[error("Failed to parse as Struct name={0}, {1}")]
40    StructParseError(String, &'static str),
41    #[error("Failed to parse as Custom Type {0}")]
42    CustomParseError(String),
43    #[error("No column found: name={0}")]
44    NoColumnFound(String),
45    #[error("invalid column index: index={0}, length={1}")]
46    InvalidColumnIndex(usize, usize),
47    #[error("invalid struct column index: index={0}")]
48    InvalidStructColumnIndex(usize),
49    #[error("No column found in struct: name={0}")]
50    NoColumnFoundInStruct(String),
51    #[error("Failed to parse as BigDecimal field={0}")]
52    BigDecimalParseError(String, #[source] ParseBigDecimalError),
53    #[error("Failed to parse as Prost Timestamp field={0}")]
54    ProstTimestampParseError(String, #[source] ::prost_types::TimestampError),
55}
56
57impl Row {
58    pub fn new(index: Arc<HashMap<String, usize>>, fields: Arc<Vec<Field>>, values: Vec<Value>) -> Row {
59        Row { index, fields, values }
60    }
61
62    pub fn column<T>(&self, column_index: usize) -> Result<T, Error>
63    where
64        T: TryFromValue,
65    {
66        column(&self.values, &self.fields, column_index)
67    }
68
69    pub fn column_by_name<T>(&self, column_name: &str) -> Result<T, Error>
70    where
71        T: TryFromValue,
72    {
73        self.column(index(&self.index, column_name)?)
74    }
75}
76
77//don't use TryFrom trait to avoid the conflict
78//https://github.com/rust-lang/rust/issues/50133
79pub trait TryFromValue: Sized {
80    fn try_from(value: &Value, field: &Field) -> Result<Self, Error>;
81}
82
83pub trait TryFromStruct: Sized {
84    fn try_from_struct(s: Struct<'_>) -> Result<Self, Error>;
85}
86
87pub struct Struct<'a> {
88    index: HashMap<String, usize>,
89    metadata: &'a StructType,
90    list_values: Option<&'a Vec<Value>>,
91    struct_values: Option<&'a BTreeMap<String, Value>>,
92}
93
94impl<'a> Struct<'a> {
95    pub fn new(metadata: &'a StructType, item: &'a Value, field: &'a Field) -> Result<Struct<'a>, Error> {
96        let kind = as_ref(item, field)?;
97        let mut index = HashMap::new();
98        for (i, f) in metadata.fields.iter().enumerate() {
99            index.insert(f.name.to_string(), i);
100        }
101        match kind {
102            Kind::ListValue(s) => Ok(Struct {
103                metadata,
104                index,
105                list_values: Some(&s.values),
106                struct_values: None,
107            }),
108            Kind::StructValue(s) => Ok(Struct {
109                metadata,
110                index,
111                list_values: None,
112                struct_values: Some(&s.fields),
113            }),
114            _ => kind_to_error(kind, field),
115        }
116    }
117
118    pub fn column<T>(&self, column_index: usize) -> Result<T, Error>
119    where
120        T: TryFromValue,
121    {
122        match self.list_values {
123            Some(values) => column(values, &self.metadata.fields, column_index),
124            None => match self.struct_values {
125                Some(values) => {
126                    let field = &self.metadata.fields[column_index];
127                    let name = &field.name;
128                    match values.get(name) {
129                        Some(value) => T::try_from(value, field),
130                        None => Err(Error::NoColumnFoundInStruct(name.to_string())),
131                    }
132                }
133                None => Err(Error::InvalidStructColumnIndex(column_index)),
134            },
135        }
136    }
137
138    pub fn column_by_name<T>(&self, column_name: &str) -> Result<T, Error>
139    where
140        T: TryFromValue,
141    {
142        self.column(index(&self.index, column_name)?)
143    }
144}
145
146impl TryFromValue for i64 {
147    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
148        match as_ref(item, field)? {
149            Kind::StringValue(s) => s.parse().map_err(|e| Error::IntParseError(field.name.to_string(), e)),
150            v => kind_to_error(v, field),
151        }
152    }
153}
154
155impl TryFromValue for f64 {
156    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
157        match as_ref(item, field)? {
158            Kind::NumberValue(s) => Ok(*s),
159            v => kind_to_error(v, field),
160        }
161    }
162}
163
164impl TryFromValue for bool {
165    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
166        match as_ref(item, field)? {
167            Kind::BoolValue(s) => Ok(*s),
168            v => kind_to_error(v, field),
169        }
170    }
171}
172
173impl TryFromValue for OffsetDateTime {
174    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
175        match as_ref(item, field)? {
176            Kind::StringValue(s) => {
177                Ok(OffsetDateTime::parse(s, &Rfc3339).map_err(|e| Error::DateParseError(field.name.to_string(), e))?)
178            }
179            v => kind_to_error(v, field),
180        }
181    }
182}
183
184impl TryFromValue for ::prost_types::Timestamp {
185    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
186        match as_ref(item, field)? {
187            Kind::StringValue(s) => Ok(::prost_types::Timestamp::from_str(s)
188                .map_err(|e| Error::ProstTimestampParseError(field.name.to_string(), e))?),
189            v => kind_to_error(v, field),
190        }
191    }
192}
193
194impl TryFromValue for CommitTimestamp {
195    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
196        Ok(CommitTimestamp {
197            timestamp: TryFromValue::try_from(item, field)?,
198        })
199    }
200}
201
202impl TryFromValue for Date {
203    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
204        match as_ref(item, field)? {
205            Kind::StringValue(s) => Date::parse(s, format_description!("[year]-[month]-[day]"))
206                .map_err(|e| Error::DateParseError(field.name.to_string(), e)),
207            v => kind_to_error(v, field),
208        }
209    }
210}
211
212impl TryFromValue for Vec<u8> {
213    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
214        match as_ref(item, field)? {
215            Kind::StringValue(s) => BASE64_STANDARD
216                .decode(s)
217                .map_err(|e| Error::ByteParseError(field.name.to_string(), e)),
218            v => kind_to_error(v, field),
219        }
220    }
221}
222
223impl TryFromValue for BigDecimal {
224    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
225        match as_ref(item, field)? {
226            Kind::StringValue(s) => {
227                Ok(BigDecimal::from_str(s).map_err(|e| Error::BigDecimalParseError(field.name.to_string(), e))?)
228            }
229            v => kind_to_error(v, field),
230        }
231    }
232}
233
234impl TryFromValue for String {
235    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
236        match as_ref(item, field)? {
237            Kind::StringValue(s) => Ok(s.to_string()),
238            v => kind_to_error(v, field),
239        }
240    }
241}
242
243impl<T> TryFromValue for T
244where
245    T: TryFromStruct,
246{
247    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
248        let maybe_array = match field.r#type.as_ref() {
249            None => return Err(Error::StructParseError(field.name.to_string(), "field type must not be none")),
250            Some(tp) => tp.array_element_type.as_ref(),
251        };
252        let maybe_struct_type = match maybe_array {
253            None => return Err(Error::StructParseError(field.name.to_string(), "array must not be none")),
254            Some(tp) => tp.struct_type.as_ref(),
255        };
256        let struct_type = match maybe_struct_type {
257            None => {
258                return Err(Error::StructParseError(
259                    field.name.to_string(),
260                    "struct type in array must not be none ",
261                ))
262            }
263            Some(struct_type) => struct_type,
264        };
265
266        T::try_from_struct(Struct::new(struct_type, item, field)?)
267    }
268}
269
270impl<T> TryFromValue for Option<T>
271where
272    T: TryFromValue,
273{
274    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
275        match as_ref(item, field)? {
276            Kind::NullValue(_i) => Ok(None),
277            _ => Ok(Some(T::try_from(item, field)?)),
278        }
279    }
280}
281
282impl<T> TryFromValue for Vec<T>
283where
284    T: TryFromValue,
285{
286    fn try_from(item: &Value, field: &Field) -> Result<Self, Error> {
287        match as_ref(item, field)? {
288            Kind::ListValue(s) => s.values.iter().map(|v| T::try_from(v, field)).collect(),
289            v => kind_to_error(v, field),
290        }
291    }
292}
293
294fn index(index: &HashMap<String, usize>, column_name: &str) -> Result<usize, Error> {
295    match index.get(column_name) {
296        Some(column_index) => Ok(*column_index),
297        None => Err(Error::NoColumnFound(column_name.to_string())),
298    }
299}
300
301fn column<T>(values: &[Value], fields: &[Field], column_index: usize) -> Result<T, Error>
302where
303    T: TryFromValue,
304{
305    if values.len() <= column_index {
306        return Err(Error::InvalidColumnIndex(column_index, values.len()));
307    }
308    let value = &values[column_index];
309    T::try_from(value, &fields[column_index])
310}
311
312pub fn as_ref<'a>(item: &'a Value, field: &'a Field) -> Result<&'a Kind, Error> {
313    match item.kind.as_ref() {
314        Some(v) => Ok(v),
315        None => Err(Error::NoKind(field.name.to_string())),
316    }
317}
318
319pub fn kind_to_error<'a, T>(v: &'a value::Kind, field: &'a Field) -> Result<T, Error> {
320    let actual = match v {
321        Kind::StringValue(_s) => "StringValue".to_string(),
322        Kind::BoolValue(_s) => "BoolValue".to_string(),
323        Kind::NumberValue(_s) => "NumberValue".to_string(),
324        Kind::ListValue(_s) => "ListValue".to_string(),
325        Kind::StructValue(_s) => "StructValue".to_string(),
326        _ => "unknown".to_string(),
327    };
328    Err(Error::KindMismatch(field.name.to_string(), actual))
329}
330
331#[cfg(test)]
332mod tests {
333    use std::collections::HashMap;
334    use std::ops::Add;
335    use std::str::FromStr;
336    use std::sync::Arc;
337
338    use prost_types::{Timestamp, Value};
339    use time::OffsetDateTime;
340
341    use google_cloud_googleapis::spanner::v1::struct_type::Field;
342
343    use crate::bigdecimal::{BigDecimal, FromPrimitive, ToPrimitive, Zero};
344    use crate::row::{Error, Row, Struct as RowStruct, TryFromStruct};
345    use crate::statement::{Kinds, ToKind, ToStruct, Types};
346    use crate::value::CommitTimestamp;
347
348    struct TestStruct {
349        pub struct_field: String,
350        pub struct_field_time: OffsetDateTime,
351        pub commit_timestamp: CommitTimestamp,
352        pub big_decimal: BigDecimal,
353        pub prost_timestamp: Timestamp,
354    }
355
356    impl TryFromStruct for TestStruct {
357        fn try_from_struct(s: RowStruct<'_>) -> Result<Self, Error> {
358            Ok(TestStruct {
359                struct_field: s.column_by_name("struct_field")?,
360                struct_field_time: s.column_by_name("struct_field_time")?,
361                commit_timestamp: s.column_by_name("commit_timestamp")?,
362                big_decimal: s.column_by_name("big_decimal")?,
363                prost_timestamp: s.column_by_name("prost_timestamp")?,
364            })
365        }
366    }
367
368    impl ToStruct for TestStruct {
369        fn to_kinds(&self) -> Kinds {
370            vec![
371                ("struct_field", self.struct_field.to_kind()),
372                ("struct_field_time", self.struct_field_time.to_kind()),
373                // value from DB is timestamp. it's not string 'spanner.commit_timestamp()'.
374                ("commit_timestamp", OffsetDateTime::from(self.commit_timestamp).to_kind()),
375                ("big_decimal", self.big_decimal.to_kind()),
376                ("prost_timestamp", self.prost_timestamp.to_kind()),
377            ]
378        }
379
380        fn get_types() -> Types {
381            vec![
382                ("struct_field", String::get_type()),
383                ("struct_field_time", OffsetDateTime::get_type()),
384                ("commit_timestamp", CommitTimestamp::get_type()),
385                ("big_decimal", BigDecimal::get_type()),
386                ("prost_timestamp", Timestamp::get_type()),
387            ]
388        }
389    }
390
391    #[test]
392    fn test_try_from() {
393        let mut index = HashMap::new();
394        index.insert("value".to_string(), 0);
395        index.insert("array".to_string(), 1);
396        index.insert("struct".to_string(), 2);
397        index.insert("decimal".to_string(), 3);
398        index.insert("timestamp".to_string(), 4);
399
400        let now = OffsetDateTime::now_utc();
401        let row = Row {
402            index: Arc::new(index),
403            fields: Arc::new(vec![
404                Field {
405                    name: "value".to_string(),
406                    r#type: Some(String::get_type()),
407                },
408                Field {
409                    name: "array".to_string(),
410                    r#type: Some(Vec::<i64>::get_type()),
411                },
412                Field {
413                    name: "struct".to_string(),
414                    r#type: Some(Vec::<TestStruct>::get_type()),
415                },
416                Field {
417                    name: "decimal".to_string(),
418                    r#type: Some(BigDecimal::get_type()),
419                },
420                Field {
421                    name: "timestamp".to_string(),
422                    r#type: Some(Timestamp::get_type()),
423                },
424            ]),
425            values: vec![
426                Value {
427                    kind: Some("aaa".to_kind()),
428                },
429                Value {
430                    kind: Some(vec![10_i64, 100_i64].to_kind()),
431                },
432                // https://cloud.google.com/spanner/docs/query-syntax?hl=ja#using_structs_with_select
433                // SELECT ARRAY(SELECT AS STRUCT * FROM TestStruct LIMIT 2) as struct
434                Value {
435                    kind: Some(
436                        vec![
437                            TestStruct {
438                                struct_field: "aaa".to_string(),
439                                struct_field_time: now,
440                                commit_timestamp: CommitTimestamp { timestamp: now },
441                                big_decimal: BigDecimal::from_str("-99999999999999999999999999999.999999999").unwrap(),
442                                prost_timestamp: Timestamp::from_str("2024-01-01T01:13:45Z").unwrap(),
443                            },
444                            TestStruct {
445                                struct_field: "bbb".to_string(),
446                                struct_field_time: now,
447                                commit_timestamp: CommitTimestamp { timestamp: now },
448                                big_decimal: BigDecimal::from_str("99999999999999999999999999999.999999999").unwrap(),
449                                prost_timestamp: Timestamp::from_str("2027-02-19T07:23:59Z").unwrap(),
450                            },
451                        ]
452                        .to_kind(),
453                    ),
454                },
455                Value {
456                    kind: Some(BigDecimal::from_f64(100.999999999999).unwrap().to_kind()),
457                },
458                Value {
459                    kind: Some(Timestamp::from_str("1999-12-31T23:59:59Z").unwrap().to_kind()),
460                },
461            ],
462        };
463
464        let value = row.column_by_name::<String>("value").unwrap();
465        let array = row.column_by_name::<Vec<i64>>("array").unwrap();
466        let struct_data = row.column_by_name::<Vec<TestStruct>>("struct").unwrap();
467        let decimal = row.column_by_name::<BigDecimal>("decimal").unwrap();
468        let ts = row.column_by_name::<Timestamp>("timestamp").unwrap();
469        assert_eq!(value, "aaa");
470        assert_eq!(array[0], 10);
471        assert_eq!(array[1], 100);
472        assert_eq!(decimal.to_f64().unwrap(), 100.999999999999);
473        assert_eq!(format!("{ts:}"), "1999-12-31T23:59:59Z");
474        assert_eq!(struct_data[0].struct_field, "aaa");
475        assert_eq!(struct_data[0].struct_field_time, now);
476        assert_eq!(
477            struct_data[0].big_decimal,
478            BigDecimal::from_str("-99999999999999999999999999999.999999999").unwrap()
479        );
480        assert_eq!(format!("{}", struct_data[0].prost_timestamp), "2024-01-01T01:13:45Z");
481        assert_eq!(struct_data[1].struct_field, "bbb");
482        assert_eq!(struct_data[1].struct_field_time, now);
483        assert_eq!(struct_data[1].commit_timestamp.timestamp, now);
484        assert_eq!(
485            struct_data[1].big_decimal,
486            BigDecimal::from_str("99999999999999999999999999999.999999999").unwrap()
487        );
488        assert_eq!(
489            struct_data[1].big_decimal.clone().add(&struct_data[0].big_decimal),
490            BigDecimal::zero()
491        );
492        assert_eq!(format!("{}", struct_data[1].prost_timestamp), "2027-02-19T07:23:59Z");
493    }
494}