From bed1e6851b2ae01cdb015f0b40827dcf3492e75c Mon Sep 17 00:00:00 2001 From: Kyle Derr Date: Tue, 30 Apr 2024 10:06:34 -0500 Subject: [PATCH] Add support for timezone-aware time types. * Add the "raw" types for time with time zone, timestamp with time zone * Add the "full" types, including those that parse the timestamp precision argument. * Ensure that these types traverse in and out of the Row type. * Support deserialization into the DateTime type (because unless you specifically ask Trino, it will return a timestamp in UTC). --- src/types/date_time.rs | 29 +++++++++++++++----- src/types/mod.rs | 61 +++++++++++++++++++++++++++++++++++++----- 2 files changed, 77 insertions(+), 13 deletions(-) diff --git a/src/types/date_time.rs b/src/types/date_time.rs index fcb516a..913c9c1 100644 --- a/src/types/date_time.rs +++ b/src/types/date_time.rs @@ -1,4 +1,5 @@ use std::fmt; +use chrono::{DateTime, Utc}; use chrono::naive::{NaiveDate, NaiveDateTime, NaiveTime}; use serde::de::{self, DeserializeSeed, Deserializer, Visitor}; @@ -6,7 +7,7 @@ use serde::de::{self, DeserializeSeed, Deserializer, Visitor}; use super::{Context, Presto, PrestoTy}; macro_rules! gen_date_time { - ($ty:ty, $seed:ident, $pty:expr, $format:expr, $empty:expr, $expect:expr) => { + ($ty:ty, $parse:expr, $seed:ident, $pty:expr, $format:expr, $empty:expr, $expect:expr, $tzmap:expr) => { impl Presto for $ty { type ValueType<'a> = String; type Seed<'a, 'de> = $seed; @@ -41,9 +42,9 @@ macro_rules! gen_date_time { where E: de::Error, { - <$ty>::parse_from_str(v, $format).map_err(|e| { + $parse(v, $format).map_err(|e| { de::Error::custom(format!("deserialize {} failed, reason: {}", $expect, e)) - }) + }).map($tzmap) } } @@ -61,14 +62,17 @@ macro_rules! gen_date_time { gen_date_time!( NaiveDate, + NaiveDate::parse_from_str, NaiveDateSeed, PrestoTy::Date, "%Y-%m-%d", NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(), - "naive date" + "naive date", + |t| t ); gen_date_time!( NaiveDateTime, + NaiveDateTime::parse_from_str, NaiveDateTimeSeed, PrestoTy::Timestamp, "%Y-%m-%d %H:%M:%S%.3f", @@ -76,13 +80,26 @@ gen_date_time!( .unwrap() .and_hms_opt(0, 0, 0) .unwrap(), - "naive time" + "naive time", + |t| t ); gen_date_time!( NaiveTime, + NaiveTime::parse_from_str, NaiveTimeSeed, PrestoTy::Time, "%H:%M:%S%.3f", NaiveTime::from_hms_opt(0, 0, 0).unwrap(), - "naive date time" + "naive date time", + |t| t ); +gen_date_time!( + DateTime, + NaiveDateTime::parse_from_str, + TimestampWithTimeZoneSeed, + PrestoTy::TimestampWithTimeZone, + "%Y-%m-%d %H:%M:%S%.3f %Z", + DateTime::default(), + "date time with time zone", + |t| t.and_utc() +); \ No newline at end of file diff --git a/src/types/mod.rs b/src/types/mod.rs index 0b96bcc..454efc1 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -50,6 +50,7 @@ use crate::{ ClientTypeSignatureParameter, Column, NamedTypeSignature, RawPrestoTy, RowFieldName, TypeSignature, }; +use crate::PrestoTy::{TimestampWithTimeZone, TimeWithTimeZone}; //TODO: refine it #[derive(Display, Debug)] @@ -125,8 +126,10 @@ fn extract(target: &PrestoTy, provided: &PrestoTy) -> Result extract(ty, provided), (Boolean, Boolean) => Ok(vec![]), (Date, Date) => Ok(vec![]), - (Time, Time) => Ok(vec![]), - (Timestamp, Timestamp) => Ok(vec![]), + (Time, Time | TimeWithPrecision(_)) => Ok(vec![]), + (TimeWithTimeZone, TimeWithTimeZone) => {Ok(vec![])}, + (Timestamp, Timestamp | TimestampWithPrecision(_)) => Ok(vec![]), + (TimestampWithTimeZone, TimestampWithTimeZone | TimestampWithTimeZoneWithPrecision(_)) => Ok(vec![]), (IntervalYearToMonth, IntervalYearToMonth) => Ok(vec![]), (IntervalDayToSecond, IntervalDayToSecond) => Ok(vec![]), (PrestoInt(_), PrestoInt(_)) => Ok(vec![]), @@ -171,14 +174,18 @@ fn extract(target: &PrestoTy, provided: &PrestoTy) -> Result PrestoTy::IpAddress, - RawPrestoTy::Uuid => PrestoTy::Uuid, + RawPrestoTy::TimestampWithTimeZone => { + match sig.arguments.as_slice() { + [] => PrestoTy::TimestampWithTimeZone, + [ClientTypeSignatureParameter::LongLiteral(s)] => { + PrestoTy::TimestampWithTimeZoneWithPrecision(*s as usize) + } + _ => return Err(Error::InvalidTypeSignature), + } + } + RawPrestoTy::TimeWithTimeZone if sig.arguments.len() == 1 => { + match sig.arguments.as_slice() { + [] => PrestoTy::TimeWithTimeZone, + [ClientTypeSignatureParameter::LongLiteral(s)] => { + PrestoTy::TimeWithPrecision(*s as usize) + } + _ => return Err(Error::InvalidTypeSignature), + } + } _ => return Err(Error::InvalidTypeSignature), }; @@ -337,7 +360,12 @@ impl PrestoTy { ], Date => vec![], Time => vec![], + TimeWithTimeZone => vec![], + TimeWithPrecision(s) => vec![ClientTypeSignatureParameter::LongLiteral(s as u64)], Timestamp => vec![], + TimestampWithPrecision(s) => vec![ClientTypeSignatureParameter::LongLiteral(s as u64)], + TimestampWithTimeZone => vec![], + TimestampWithTimeZoneWithPrecision(s) => vec![ClientTypeSignatureParameter::LongLiteral(s as u64)], IntervalYearToMonth => vec![], IntervalDayToSecond => vec![], Option(t) => return t.into_type_signature(), @@ -381,7 +409,24 @@ impl PrestoTy { Option(t) => t.full_type(), Date => RawPrestoTy::Date.to_str().into(), Time => RawPrestoTy::Time.to_str().into(), + TimeWithPrecision(p) => format!( + "{}({})", + RawPrestoTy::Time.to_str(), + p.to_string() + ).into(), + TimestampWithPrecision(p) => format!( + "{}({})", + RawPrestoTy::Timestamp.to_str(), + p.to_string() + ).into(), + TimestampWithTimeZoneWithPrecision(p) => format!( + "{}({})", + RawPrestoTy::TimestampWithTimeZone.to_str(), + p.to_string() + ).into(), + TimeWithTimeZone => RawPrestoTy::TimeWithTimeZone.to_str().into(), Timestamp => RawPrestoTy::Timestamp.to_str().into(), + TimestampWithTimeZone => RawPrestoTy::TimestampWithTimeZone.to_str().into(), IntervalYearToMonth => RawPrestoTy::IntervalYearToMonth.to_str().into(), IntervalDayToSecond => RawPrestoTy::IntervalDayToSecond.to_str().into(), Boolean => RawPrestoTy::Boolean.to_str().into(), @@ -421,8 +466,10 @@ impl PrestoTy { match self { Unknown => RawPrestoTy::Unknown, Date => RawPrestoTy::Date, - Time => RawPrestoTy::Time, - Timestamp => RawPrestoTy::Timestamp, + Time | TimeWithPrecision(_) => RawPrestoTy::Time, + TimeWithTimeZone => RawPrestoTy::TimeWithTimeZone, + Timestamp | TimestampWithPrecision(_) => RawPrestoTy::Timestamp, + TimestampWithTimeZone | TimestampWithTimeZoneWithPrecision(_) => RawPrestoTy::TimestampWithTimeZone, IntervalYearToMonth => RawPrestoTy::IntervalYearToMonth, IntervalDayToSecond => RawPrestoTy::IntervalDayToSecond, Decimal(_, _) => RawPrestoTy::Decimal,