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
77pub 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 ("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 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}