chore: extract timestamp value to pg-srv

Author: ovr

rust/cubesql/Cargo.lock

@@ -1,12 +1,4 @@
-use chrono::{
-    format::{
-        Fixed, Item,
-        Numeric::{Day, Hour, Minute, Month, Second, Year},
-        Pad::Zero,
-    },
-    prelude::*,
-};
-use chrono_tz::Tz;
+use chrono::prelude::*;
 use comfy_table::{Cell, Table};
 use datafusion::arrow::{
     array::{
@@ -18,15 +10,13 @@ use datafusion::arrow::{
     },
     datatypes::{DataType, IntervalUnit, Schema, TimeUnit},
     record_batch::RecordBatch,
-    temporal_conversions,
 };
-use pg_srv::IntervalValue;
+// Type aliases for compatibility - actual implementations are in pg-srv
+pub type IntervalValue = pg_srv::IntervalValue;
+pub type TimestampValue = pg_srv::TimestampValue;
 use rust_decimal::prelude::*;
 use serde::{Serialize, Serializer};
-use std::{
-    fmt::{self, Debug, Formatter},
-    io,
-};
+use std::fmt::Debug;
 
 use super::{ColumnFlags, ColumnType};
 use crate::CubeError;
@@ -207,83 +197,6 @@ impl DataFrame {
     }
 }
 
-#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
-pub struct TimestampValue {
-    unix_nano: i64,
-    tz: Option<String>,
-}
-
-impl TimestampValue {
-    pub fn new(mut unix_nano: i64, tz: Option<String>) -> TimestampValue {
-        // This is a hack to workaround a mismatch between on-disk and in-memory representations.
-        // We use millisecond precision on-disk.
-        unix_nano -= unix_nano % 1000;
-        TimestampValue { unix_nano, tz }
-    }
-
-    pub fn to_naive_datetime(&self) -> NaiveDateTime {
-        assert!(self.tz.is_none());
-
-        temporal_conversions::timestamp_ns_to_datetime(self.unix_nano)
-    }
-
-    pub fn to_fixed_datetime(&self) -> io::Result<DateTime<Tz>> {
-        assert!(self.tz.is_some());
-
-        let tz = self
-            .tz
-            .as_ref()
-            .unwrap()
-            .parse::<Tz>()
-            .map_err(|err| io::Error::new(io::ErrorKind::Other, err.to_string()))?;
-
-        let ndt = temporal_conversions::timestamp_ns_to_datetime(self.unix_nano);
-        Ok(tz.from_utc_datetime(&ndt))
-    }
-
-    pub fn tz_ref(&self) -> &Option<String> {
-        &self.tz
-    }
-
-    pub fn get_time_stamp(&self) -> i64 {
-        self.unix_nano
-    }
-}
-
-impl Debug for TimestampValue {
-    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        f.debug_struct("TimestampValue")
-            .field("unix_nano", &self.unix_nano)
-            .field("tz", &self.tz)
-            .field("str", &self.to_string())
-            .finish()
-    }
-}
-
-impl ToString for TimestampValue {
-    fn to_string(&self) -> String {
-        Utc.timestamp_nanos(self.unix_nano)
-            .format_with_items(
-                [
-                    Item::Numeric(Year, Zero),
-                    Item::Literal("-"),
-                    Item::Numeric(Month, Zero),
-                    Item::Literal("-"),
-                    Item::Numeric(Day, Zero),
-                    Item::Literal("T"),
-                    Item::Numeric(Hour, Zero),
-                    Item::Literal(":"),
-                    Item::Numeric(Minute, Zero),
-                    Item::Literal(":"),
-                    Item::Numeric(Second, Zero),
-                    Item::Fixed(Fixed::Nanosecond3),
-                ]
-                .iter(),
-            )
-            .to_string()
-    }
-}
-
 #[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
 pub struct Decimal128Value {
     n: i128,

rust/cubesql/cubesql/src/sql/dataframe.rs

@@ -1,16 +1,8 @@
 use crate::sql::{
-    dataframe::{Decimal128Value, ListValue, TimestampValue},
+    dataframe::{Decimal128Value, ListValue},
     df_type_to_pg_tid,
 };
 use bytes::{BufMut, BytesMut};
-use chrono::{
-    format::{
-        Fixed, Item,
-        Numeric::{Day, Hour, Minute, Month, Second, Year},
-        Pad::Zero,
-    },
-    prelude::*,
-};
 use datafusion::arrow::{
     array::{
         Array, BooleanArray, Float16Array, Float32Array, Float64Array, Int16Array, Int32Array,
@@ -24,87 +16,7 @@ use pg_srv::{
     PgTypeId, ProtocolError, ToProtocolValue,
 };
 use postgres_types::{ToSql, Type};
-use std::{convert::TryFrom, io, io::Error};
-
-// POSTGRES_EPOCH_JDATE
-fn pg_base_date_epoch() -> NaiveDateTime {
-    NaiveDate::from_ymd_opt(2000, 1, 1)
-        .unwrap()
-        .and_hms_opt(0, 0, 0)
-        .unwrap()
-}
-
-impl ToProtocolValue for TimestampValue {
-    fn to_text(&self, buf: &mut BytesMut) -> Result<(), ProtocolError> {
-        let ndt = match self.tz_ref() {
-            None => self.to_naive_datetime(),
-            Some(_) => self.to_fixed_datetime()?.naive_utc(),
-        };
-
-        // 2022-04-25 15:36:49.39705+00
-        let as_str = ndt
-            .format_with_items(
-                [
-                    Item::Numeric(Year, Zero),
-                    Item::Literal("-"),
-                    Item::Numeric(Month, Zero),
-                    Item::Literal("-"),
-                    Item::Numeric(Day, Zero),
-                    Item::Literal(" "),
-                    Item::Numeric(Hour, Zero),
-                    Item::Literal(":"),
-                    Item::Numeric(Minute, Zero),
-                    Item::Literal(":"),
-                    Item::Numeric(Second, Zero),
-                    Item::Fixed(Fixed::Nanosecond6),
-                ]
-                .iter(),
-            )
-            .to_string();
-
-        match self.tz_ref() {
-            None => as_str.to_text(buf),
-            Some(_) => (as_str + "+00").to_text(buf),
-        }
-    }
-
-    fn to_binary(&self, buf: &mut BytesMut) -> Result<(), ProtocolError> {
-        match self.tz_ref() {
-            None => {
-                let seconds = self
-                    .to_naive_datetime()
-                    .signed_duration_since(pg_base_date_epoch())
-                    .num_microseconds()
-                    .ok_or(Error::new(
-                        io::ErrorKind::Other,
-                        "Unable to extract number of seconds from timestamp",
-                    ))?;
-
-                buf.put_i32(8_i32);
-                buf.put_i64(seconds)
-            }
-            Some(tz) => {
-                let seconds = self
-                    .to_fixed_datetime()?
-                    .naive_utc()
-                    .signed_duration_since(pg_base_date_epoch())
-                    .num_microseconds()
-                    .ok_or(Error::new(
-                        io::ErrorKind::Other,
-                        format!(
-                            "Unable to extract number of seconds from timestamp with tz: {}",
-                            tz
-                        ),
-                    ))?;
-
-                buf.put_i32(8_i32);
-                buf.put_i64(seconds)
-            }
-        };
-
-        Ok(())
-    }
-}
+use std::convert::TryFrom;
 
 /// https://github.com/postgres/postgres/blob/REL_14_4/src/backend/utils/adt/numeric.c#L1022
 impl ToProtocolValue for Decimal128Value {
@@ -324,13 +236,13 @@ impl Serialize for BatchWriter {
 #[cfg(test)]
 mod tests {
     use crate::sql::{
-        dataframe::{Decimal128Value, ListValue, TimestampValue},
+        dataframe::{Decimal128Value, ListValue},
         shim::ConnectionError,
         writer::{BatchWriter, ToProtocolValue},
     };
     use bytes::BytesMut;
     use datafusion::arrow::array::{ArrayRef, Int64Builder};
-    use pg_srv::{buffer, protocol::Format};
+    use pg_srv::{buffer, protocol::Format, TimestampValue};
     use std::{io::Cursor, sync::Arc};
 
     fn assert_text_encode<T: ToProtocolValue>(value: T, expected: &[u8]) {

rust/cubesql/cubesql/src/sql/postgres/writer.rs

@@ -22,6 +22,7 @@ thiserror = "2"
 chrono = { version = "0.4", package = "chrono", default-features = false, features = [
     "clock",
 ], optional = true }
+chrono-tz = "0.8"
 
 [dev-dependencies]
 hex = "0.4.3"

rust/cubesql/pg-srv/Cargo.toml

@@ -146,7 +146,6 @@ impl ToProtocolValue for NaiveDate {
     }
 }
 
-
 #[cfg(test)]
 mod tests {
     use crate::*;
@@ -182,5 +181,4 @@ mod tests {
 
         Ok(())
     }
-
 }

rust/cubesql/pg-srv/src/encoding.rs

@@ -204,4 +204,4 @@ mod tests {
 
         Ok(())
     }
-}
+}

rust/cubesql/pg-srv/src/interval.rs

@@ -10,13 +10,15 @@ pub mod extended;
 pub mod interval;
 pub mod pg_type;
 pub mod protocol;
+pub mod timestamp;
 
 pub use buffer::*;
 pub use decoding::*;
 pub use encoding::*;
 pub use extended::*;
 pub use interval::*;
 pub use pg_type::*;
+pub use timestamp::*;
 
 use std::{backtrace::Backtrace, fmt::Formatter};