From fa40e9e55f65721514ec0f06c2d576b5c7f8983f Mon Sep 17 00:00:00 2001
From: Ryan Leckey <ryan@launchbadge.com>
Date: Fri, 17 Jul 2020 03:22:33 -0700
Subject: [PATCH] implement support for postgres TIMETZ type

Co-authored-by: Julius de Bruijn <julius+github@nauk.io>
---
 .../src/postgres/types/chrono/datetime.rs     |  11 +-
 sqlx-core/src/postgres/types/mod.rs           |  10 +
 sqlx-core/src/postgres/types/time_tz.rs       | 183 ++++++++++++++++++
 sqlx-macros/src/database/postgres.rs          |   9 +-
 tests/postgres/types.rs                       |  31 ++-
 5 files changed, 238 insertions(+), 6 deletions(-)
 create mode 100644 sqlx-core/src/postgres/types/time_tz.rs

diff --git a/sqlx-core/src/postgres/types/chrono/datetime.rs b/sqlx-core/src/postgres/types/chrono/datetime.rs
index ce97c700..a8712e7b 100644
--- a/sqlx-core/src/postgres/types/chrono/datetime.rs
+++ b/sqlx-core/src/postgres/types/chrono/datetime.rs
@@ -3,7 +3,9 @@ use crate::encode::{Encode, IsNull};
 use crate::error::BoxDynError;
 use crate::postgres::{PgArgumentBuffer, PgTypeInfo, PgValueFormat, PgValueRef, Postgres};
 use crate::types::Type;
-use chrono::{DateTime, Duration, Local, NaiveDate, NaiveDateTime, TimeZone, Utc};
+use chrono::{
+    DateTime, Duration, FixedOffset, Local, NaiveDate, NaiveDateTime, Offset, TimeZone, Utc,
+};
 use std::mem;
 
 impl Type<Postgres> for NaiveDateTime {
@@ -110,3 +112,10 @@ impl<'r> Decode<'r, Postgres> for DateTime<Utc> {
         Ok(Utc.from_utc_datetime(&naive))
     }
 }
+
+impl<'r> Decode<'r, Postgres> for DateTime<FixedOffset> {
+    fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
+        let naive = <NaiveDateTime as Decode<Postgres>>::decode(value)?;
+        Ok(Utc.fix().from_utc_datetime(&naive))
+    }
+}
diff --git a/sqlx-core/src/postgres/types/mod.rs b/sqlx-core/src/postgres/types/mod.rs
index 4a808239..811ee847 100644
--- a/sqlx-core/src/postgres/types/mod.rs
+++ b/sqlx-core/src/postgres/types/mod.rs
@@ -46,6 +46,7 @@
 //! | `chrono::NaiveDateTime`               | TIMESTAMP                                            |
 //! | `chrono::NaiveDate`                   | DATE                                                 |
 //! | `chrono::NaiveTime`                   | TIME                                                 |
+//! | [`PgTimeTz`]                          | TIMETZ                                               |
 //!
 //! ### [`time`](https://crates.io/crates/time)
 //!
@@ -57,6 +58,9 @@
 //! | `time::OffsetDateTime`                | TIMESTAMPTZ                                          |
 //! | `time::Date`                          | DATE                                                 |
 //! | `time::Time`                          | TIME                                                 |
+//! | [`PgTimeTz`]                          | TIMETZ                                               |
+//!
+//! [`PgTimeTz`]: struct.PgTimeTz.html
 //!
 //! ### [`uuid`](https://crates.io/crates/uuid)
 //!
@@ -162,6 +166,9 @@ mod str;
 mod tuple;
 mod void;
 
+#[cfg(any(feature = "chrono", feature = "time"))]
+mod time_tz;
+
 #[cfg(feature = "bigdecimal")]
 mod bigdecimal;
 
@@ -190,6 +197,9 @@ pub use interval::PgInterval;
 pub use money::PgMoney;
 pub use range::PgRange;
 
+#[cfg(any(feature = "chrono", feature = "time"))]
+pub use time_tz::PgTimeTz;
+
 // used in derive(Type) for `struct`
 // but the interface is not considered part of the public API
 #[doc(hidden)]
diff --git a/sqlx-core/src/postgres/types/time_tz.rs b/sqlx-core/src/postgres/types/time_tz.rs
new file mode 100644
index 00000000..0129f482
--- /dev/null
+++ b/sqlx-core/src/postgres/types/time_tz.rs
@@ -0,0 +1,183 @@
+use crate::decode::Decode;
+use crate::encode::{Encode, IsNull};
+use crate::error::BoxDynError;
+use crate::postgres::{PgArgumentBuffer, PgTypeInfo, PgValueFormat, PgValueRef, Postgres};
+use crate::types::Type;
+use byteorder::{BigEndian, ReadBytesExt};
+use std::io::Cursor;
+use std::mem;
+
+#[cfg(feature = "time")]
+type DefaultTime = ::time::Time;
+
+#[cfg(all(not(feature = "time"), feature = "chrono"))]
+type DefaultTime = ::chrono::NaiveTime;
+
+#[cfg(feature = "time")]
+type DefaultOffset = ::time::UtcOffset;
+
+#[cfg(all(not(feature = "time"), feature = "chrono"))]
+type DefaultOffset = ::chrono::FixedOffset;
+
+/// Represents a moment of time, in a specified timezone.
+///
+/// # Warning
+///
+/// `PgTimeTz` provides `TIMETZ` and is supported only for reading from legacy databases.
+/// [PostgreSQL recommends] to use `TIMESTAMPTZ` instead.
+///
+/// [PostgreSQL recommends]: https://wiki.postgresql.org/wiki/Don't_Do_This#Don.27t_use_timetz
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub struct PgTimeTz<Time = DefaultTime, Offset = DefaultOffset> {
+    pub time: Time,
+    pub offset: Offset,
+}
+
+impl<Time, Offset> Type<Postgres> for [PgTimeTz<Time, Offset>]
+where
+    PgTimeTz<Time, Offset>: Type<Postgres>,
+{
+    fn type_info() -> PgTypeInfo {
+        PgTypeInfo::TIMETZ_ARRAY
+    }
+}
+
+impl<Time, Offset> Type<Postgres> for Vec<PgTimeTz<Time, Offset>>
+where
+    PgTimeTz<Time, Offset>: Type<Postgres>,
+{
+    fn type_info() -> PgTypeInfo {
+        PgTypeInfo::TIMETZ_ARRAY
+    }
+}
+
+#[cfg(feature = "chrono")]
+mod chrono {
+    use super::*;
+    use ::chrono::{DateTime, Duration, FixedOffset, NaiveTime};
+
+    impl Type<Postgres> for PgTimeTz<NaiveTime, FixedOffset> {
+        fn type_info() -> PgTypeInfo {
+            PgTypeInfo::TIMETZ
+        }
+    }
+
+    impl Encode<'_, Postgres> for PgTimeTz<NaiveTime, FixedOffset> {
+        fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> IsNull {
+            let _ = <NaiveTime as Encode<'_, Postgres>>::encode(self.time, buf);
+            let _ = <i32 as Encode<'_, Postgres>>::encode(self.offset.utc_minus_local(), buf);
+
+            IsNull::No
+        }
+
+        fn size_hint(&self) -> usize {
+            mem::size_of::<i64>() + mem::size_of::<i32>()
+        }
+    }
+
+    impl<'r> Decode<'r, Postgres> for PgTimeTz<NaiveTime, FixedOffset> {
+        fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
+            match value.format() {
+                PgValueFormat::Binary => {
+                    let mut buf = Cursor::new(value.as_bytes()?);
+
+                    // TIME is encoded as the microseconds since midnight
+                    let us = buf.read_i64::<BigEndian>()?;
+                    let time = NaiveTime::from_hms(0, 0, 0) + Duration::microseconds(us);
+
+                    // OFFSET is encoded as seconds from UTC
+                    let seconds = buf.read_i32::<BigEndian>()?;
+
+                    Ok(PgTimeTz {
+                        time,
+                        offset: FixedOffset::west(seconds),
+                    })
+                }
+
+                PgValueFormat::Text => {
+                    let s = value.as_str()?;
+
+                    let mut tmp = String::with_capacity(11 + s.len());
+                    tmp.push_str("2001-07-08 ");
+                    tmp.push_str(s);
+
+                    let dt = 'out: loop {
+                        let mut err = None;
+
+                        for fmt in &["%Y-%m-%d %H:%M:%S%.f%#z", "%Y-%m-%d %H:%M:%S%.f"] {
+                            match DateTime::parse_from_str(&tmp, fmt) {
+                                Ok(dt) => {
+                                    break 'out dt;
+                                }
+
+                                Err(error) => {
+                                    err = Some(error);
+                                }
+                            }
+                        }
+
+                        return Err(err.unwrap().into());
+                    };
+
+                    let time = dt.time();
+                    let offset = *dt.offset();
+
+                    Ok(PgTimeTz { time, offset })
+                }
+            }
+        }
+    }
+}
+
+#[cfg(feature = "time")]
+mod time {
+    use super::*;
+    use ::time::{Duration, Time, UtcOffset};
+
+    impl Type<Postgres> for PgTimeTz<Time, UtcOffset> {
+        fn type_info() -> PgTypeInfo {
+            PgTypeInfo::TIMETZ
+        }
+    }
+
+    impl Encode<'_, Postgres> for PgTimeTz<Time, UtcOffset> {
+        fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> IsNull {
+            let _ = <Time as Encode<'_, Postgres>>::encode(self.time, buf);
+            let _ = <i32 as Encode<'_, Postgres>>::encode(-self.offset.as_seconds(), buf);
+
+            IsNull::No
+        }
+
+        fn size_hint(&self) -> usize {
+            mem::size_of::<i64>() + mem::size_of::<i32>()
+        }
+    }
+
+    impl<'r> Decode<'r, Postgres> for PgTimeTz<Time, UtcOffset> {
+        fn decode(value: PgValueRef<'r>) -> Result<Self, BoxDynError> {
+            match value.format() {
+                PgValueFormat::Binary => {
+                    let mut buf = Cursor::new(value.as_bytes()?);
+
+                    // TIME is encoded as the microseconds since midnight
+                    let us = buf.read_i64::<BigEndian>()?;
+                    let time = Time::midnight() + Duration::microseconds(us);
+
+                    // OFFSET is encoded as seconds from UTC
+                    let seconds = buf.read_i32::<BigEndian>()?;
+
+                    Ok(PgTimeTz {
+                        time,
+                        offset: UtcOffset::west_seconds(seconds as u32),
+                    })
+                }
+
+                PgValueFormat::Text => {
+                    // the `time` crate has a limited ability to parse and can't parse the
+                    // timezone format
+                    Err("reading a `TIMETZ` value in text format is not supported.".into())
+                }
+            }
+        }
+    }
+}
diff --git a/sqlx-macros/src/database/postgres.rs b/sqlx-macros/src/database/postgres.rs
index bb551e98..b8aa7e1d 100644
--- a/sqlx-macros/src/database/postgres.rs
+++ b/sqlx-macros/src/database/postgres.rs
@@ -11,9 +11,10 @@ impl_database_ext! {
         i64,
         f32,
         f64,
-
         Vec<u8> | &[u8],
 
+        sqlx::postgres::types::PgInterval,
+
         #[cfg(feature = "uuid")]
         sqlx::types::Uuid,
 
@@ -29,6 +30,9 @@ impl_database_ext! {
         #[cfg(feature = "chrono")]
         sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc> | sqlx::types::chrono::DateTime<_>,
 
+        #[cfg(feature = "chrono")]
+        sqlx::postgres::types::PgTimeTz<sqlx::types::chrono::NaiveTime, sqlx::types::chrono::FixedOffset>,
+
         #[cfg(feature = "time")]
         sqlx::types::time::Time,
 
@@ -41,7 +45,8 @@ impl_database_ext! {
         #[cfg(feature = "time")]
         sqlx::types::time::OffsetDateTime,
 
-        sqlx::postgres::types::PgInterval,
+        #[cfg(feature = "time")]
+        sqlx::postgres::types::PgTimeTz<sqlx::types::time::Time, sqlx::types::time::UtcOffset>,
 
         #[cfg(feature = "bigdecimal")]
         sqlx::types::BigDecimal,
diff --git a/tests/postgres/types.rs b/tests/postgres/types.rs
index b4302ee1..da1dc80b 100644
--- a/tests/postgres/types.rs
+++ b/tests/postgres/types.rs
@@ -179,7 +179,11 @@ test_type!(ipnetwork_vec<Vec<sqlx::types::ipnetwork::IpNetwork>>(Postgres,
 #[cfg(feature = "chrono")]
 mod chrono {
     use super::*;
-    use sqlx::types::chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
+    use sqlx::types::chrono::{
+        DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc,
+    };
+
+    type PgTimeTz = sqlx::postgres::types::PgTimeTz<NaiveTime, FixedOffset>;
 
     test_type!(chrono_date<NaiveDate>(Postgres,
         "DATE '2001-01-05'" == NaiveDate::from_ymd(2001, 1, 5),
@@ -199,7 +203,7 @@ mod chrono {
             == vec![NaiveDate::from_ymd(2019, 1, 2).and_hms(5, 10, 20)]
     ));
 
-    test_type!(chrono_date_time_tz<DateTime::<Utc>>(Postgres,
+    test_type!(chrono_date_time_tz_utc<DateTime::<Utc>>(Postgres,
         "TIMESTAMPTZ '2019-01-02 05:10:20.115100'"
             == DateTime::<Utc>::from_utc(
                 NaiveDate::from_ymd(2019, 1, 2).and_hms_micro(5, 10, 20, 115100),
@@ -207,6 +211,11 @@ mod chrono {
             )
     ));
 
+    test_type!(chrono_date_time_tz<DateTime::<FixedOffset>>(Postgres,
+        "TIMESTAMPTZ '2019-01-02 05:10:20.115100+06:30'"
+            == FixedOffset::east(60 * 60 * 6 + 1800).ymd(2019, 1, 2).and_hms_micro(5, 10, 20, 115100)
+    ));
+
     test_type!(chrono_date_time_tz_vec<Vec<DateTime::<Utc>>>(Postgres,
         "array['2019-01-02 05:10:20.115100']::timestamptz[]"
             == vec![
@@ -216,14 +225,23 @@ mod chrono {
                 )
             ]
     ));
+
+    test_type!(chrono_time_tz<PgTimeTz>(Postgres,
+        "TIMETZ '05:10:20.115100+00'" == PgTimeTz { time: NaiveTime::from_hms_micro(5, 10, 20, 115100), offset: FixedOffset::east(0) },
+        "TIMETZ '05:10:20.115100+06:30'" == PgTimeTz { time: NaiveTime::from_hms_micro(5, 10, 20, 115100), offset: FixedOffset::east(60 * 60 * 6 + 1800) },
+        "TIMETZ '05:10:20.115100-05'" == PgTimeTz { time: NaiveTime::from_hms_micro(5, 10, 20, 115100), offset: FixedOffset::west(60 * 60 * 5) },
+        "TIMETZ '05:10:20+02'" == PgTimeTz { time: NaiveTime::from_hms(5, 10, 20), offset: FixedOffset::east(60 * 60 * 2 )}
+    ));
 }
 
 #[cfg(feature = "time")]
 mod time_tests {
     use super::*;
-    use sqlx::types::time::{Date, OffsetDateTime, PrimitiveDateTime, Time};
+    use sqlx::types::time::{Date, OffsetDateTime, PrimitiveDateTime, Time, UtcOffset};
     use time::{date, time};
 
+    type PgTimeTz = sqlx::postgres::types::PgTimeTz<Time, UtcOffset>;
+
     test_type!(time_date<Date>(
         Postgres,
         "DATE '2001-01-05'" == date!(2001 - 1 - 5),
@@ -249,6 +267,13 @@ mod time_tests {
                 .with_time(time!(5:10:20.115100))
                 .assume_utc()
     ));
+
+    test_prepared_type!(time_time_tz<PgTimeTz>(Postgres,
+        "TIMETZ '05:10:20.115100+00'" == PgTimeTz { time: time!(5:10:20.115100), offset: UtcOffset::east_seconds(0) },
+        "TIMETZ '05:10:20.115100+06:30'" == PgTimeTz { time: time!(5:10:20.115100), offset: UtcOffset::east_seconds(60 * 60 * 6 + 1800) },
+        "TIMETZ '05:10:20.115100-05'" == PgTimeTz { time: time!(5:10:20.115100), offset: UtcOffset::west_seconds(60 * 60 * 5) },
+        "TIMETZ '05:10:20+02'" == PgTimeTz { time: time!(5:10:20), offset: UtcOffset::east_seconds(60 * 60 * 2 )}
+    ));
 }
 
 #[cfg(feature = "json")]