feat!: parsing local timestamp (#1352)

* fix: parse and display timestamp/datetime in local time zone

* fix display

* fix: unit tests

* change time zone env

* fix: remove useless code

src/common/time/src/datetime.rs

@@ -15,13 +15,13 @@
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
 
-use chrono::NaiveDateTime;
+use chrono::{Local, LocalResult, NaiveDateTime, TimeZone};
 use serde::{Deserialize, Serialize};
-use snafu::ResultExt;
 
-use crate::error::{Error, ParseDateStrSnafu, Result};
+use crate::error::{Error, InvalidDateStrSnafu, Result};
 
 const DATETIME_FORMAT: &str = "%F %T";
+const DATETIME_FORMAT_WITH_TZ: &str = "%F %T%z";
 
 /// [DateTime] represents the **seconds elapsed since "1970-01-01 00:00:00 UTC" (UNIX Epoch)**.
 #[derive(
@@ -32,7 +32,13 @@ pub struct DateTime(i64);
 impl Display for DateTime {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         if let Some(abs_time) = NaiveDateTime::from_timestamp_opt(self.0, 0) {
-            write!(f, "{}", abs_time.format(DATETIME_FORMAT))
+            write!(
+                f,
+                "{}",
+                Local {}
+                    .from_utc_datetime(&abs_time)
+                    .format(DATETIME_FORMAT_WITH_TZ)
+            )
         } else {
             write!(f, "DateTime({})", self.0)
         }
@@ -49,9 +55,21 @@ impl FromStr for DateTime {
     type Err = Error;
 
     fn from_str(s: &str) -> Result<Self> {
-        let datetime = NaiveDateTime::parse_from_str(s, DATETIME_FORMAT)
+        let local = Local {};
-            .context(ParseDateStrSnafu { raw: s })?;
+        let timestamp = if let Ok(d) = NaiveDateTime::parse_from_str(s, DATETIME_FORMAT) {
-        Ok(Self(datetime.timestamp()))
+            match local.from_local_datetime(&d) {
+                LocalResult::None => {
+                    return InvalidDateStrSnafu { raw: s }.fail();
+                }
+                LocalResult::Single(d) | LocalResult::Ambiguous(d, _) => d.naive_utc().timestamp(),
+            }
+        } else if let Ok(v) = chrono::DateTime::parse_from_str(s, DATETIME_FORMAT_WITH_TZ) {
+            v.timestamp()
+        } else {
+            return InvalidDateStrSnafu { raw: s }.fail();
+        };
+
+        Ok(Self(timestamp))
     }
 }
 
@@ -81,14 +99,16 @@ mod tests {
 
     #[test]
     pub fn test_new_date_time() {
-        assert_eq!("1970-01-01 00:00:00", DateTime::new(0).to_string());
+        std::env::set_var("TZ", "Asia/Shanghai");
-        assert_eq!("1970-01-01 00:00:01", DateTime::new(1).to_string());
+        assert_eq!("1970-01-01 08:00:00+0800", DateTime::new(0).to_string());
-        assert_eq!("1969-12-31 23:59:59", DateTime::new(-1).to_string());
+        assert_eq!("1970-01-01 08:00:01+0800", DateTime::new(1).to_string());
+        assert_eq!("1970-01-01 07:59:59+0800", DateTime::new(-1).to_string());
     }
 
     #[test]
     pub fn test_parse_from_string() {
-        let time = "1970-01-01 00:00:00";
+        std::env::set_var("TZ", "Asia/Shanghai");
+        let time = "1970-01-01 00:00:00+0800";
         let dt = DateTime::from_str(time).unwrap();
         assert_eq!(time, &dt.to_string());
     }
@@ -98,4 +118,22 @@ mod tests {
         let d: DateTime = 42.into();
         assert_eq!(42, d.val());
     }
+
+    #[test]
+    fn test_parse_local_date_time() {
+        std::env::set_var("TZ", "Asia/Shanghai");
+        assert_eq!(
+            -28800,
+            DateTime::from_str("1970-01-01 00:00:00").unwrap().val()
+        );
+        assert_eq!(0, DateTime::from_str("1970-01-01 08:00:00").unwrap().val());
+    }
+
+    #[test]
+    fn test_parse_local_date_time_with_tz() {
+        let ts = DateTime::from_str("1970-01-01 08:00:00+0000")
+            .unwrap()
+            .val();
+        assert_eq!(28800, ts);
+    }
 }

src/common/time/src/error.rs

@@ -26,6 +26,9 @@ pub enum Error {
     #[snafu(display("Failed to parse string to date, raw: {}, source: {}", raw, source))]
     ParseDateStr { raw: String, source: ParseError },
 
+    #[snafu(display("Invalid date string, raw: {}", raw))]
+    InvalidDateStr { raw: String, location: Location },
+
     #[snafu(display("Failed to parse a string into Timestamp, raw string: {}", raw))]
     ParseTimestamp { raw: String, location: Location },
 
@@ -46,7 +49,9 @@ impl ErrorExt for Error {
                 StatusCode::InvalidArguments
             }
             Error::TimestampOverflow { .. } => StatusCode::Internal,
-            Error::ArithmeticOverflow { .. } => StatusCode::InvalidArguments,
+            Error::InvalidDateStr { .. } | Error::ArithmeticOverflow { .. } => {
+                StatusCode::InvalidArguments
+            }
         }
     }
 
@@ -60,6 +65,7 @@ impl ErrorExt for Error {
             | Error::TimestampOverflow { location, .. }
             | Error::ArithmeticOverflow { location, .. } => Some(*location),
             Error::ParseDateStr { .. } => None,
+            Error::InvalidDateStr { location, .. } => Some(*location),
         }
     }
 }

src/common/time/src/timestamp.rs

@@ -164,16 +164,20 @@ impl Timestamp {
     /// Format timestamp to ISO8601 string. If the timestamp exceeds what chrono timestamp can
     /// represent, this function simply print the timestamp unit and value in plain string.
     pub fn to_iso8601_string(&self) -> String {
-        if let LocalResult::Single(datetime) = self.to_chrono_datetime() {
+        if let Some(v) = self.to_chrono_datetime() {
-            format!("{}", datetime.format("%Y-%m-%d %H:%M:%S%.f%z"))
+            let local = Local {};
+            format!(
+                "{}",
+                local.from_utc_datetime(&v).format("%Y-%m-%d %H:%M:%S%.f%z")
+            )
         } else {
             format!("[Timestamp{}: {}]", self.unit, self.value)
         }
     }
 
-    pub fn to_chrono_datetime(&self) -> LocalResult<DateTime<Utc>> {
+    pub fn to_chrono_datetime(&self) -> Option<NaiveDateTime> {
         let (sec, nsec) = self.split();
-        Utc.timestamp_opt(sec, nsec)
+        NaiveDateTime::from_timestamp_opt(sec, nsec)
     }
 }
 
@@ -636,31 +640,33 @@ mod tests {
 
     #[test]
     fn test_to_iso8601_string() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         let datetime_str = "2020-09-08 13:42:29.042+0000";
         let ts = Timestamp::from_str(datetime_str).unwrap();
-        assert_eq!(datetime_str, ts.to_iso8601_string());
+        assert_eq!("2020-09-08 21:42:29.042+0800", ts.to_iso8601_string());
 
         let ts_millis = 1668070237000;
         let ts = Timestamp::new_millisecond(ts_millis);
-        assert_eq!("2022-11-10 08:50:37+0000", ts.to_iso8601_string());
+        assert_eq!("2022-11-10 16:50:37+0800", ts.to_iso8601_string());
 
         let ts_millis = -1000;
         let ts = Timestamp::new_millisecond(ts_millis);
-        assert_eq!("1969-12-31 23:59:59+0000", ts.to_iso8601_string());
+        assert_eq!("1970-01-01 07:59:59+0800", ts.to_iso8601_string());
 
         let ts_millis = -1;
         let ts = Timestamp::new_millisecond(ts_millis);
-        assert_eq!("1969-12-31 23:59:59.999+0000", ts.to_iso8601_string());
+        assert_eq!("1970-01-01 07:59:59.999+0800", ts.to_iso8601_string());
 
         let ts_millis = -1001;
         let ts = Timestamp::new_millisecond(ts_millis);
-        assert_eq!("1969-12-31 23:59:58.999+0000", ts.to_iso8601_string());
+        assert_eq!("1970-01-01 07:59:58.999+0800", ts.to_iso8601_string());
     }
 
     #[test]
     fn test_serialize_to_json_value() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         assert_eq!(
-            "1970-01-01 00:00:01+0000",
+            "1970-01-01 08:00:01+0800",
             match serde_json::Value::from(Timestamp::new(1, TimeUnit::Second)) {
                 Value::String(s) => s,
                 _ => unreachable!(),
@@ -668,7 +674,7 @@ mod tests {
         );
 
         assert_eq!(
-            "1970-01-01 00:00:00.001+0000",
+            "1970-01-01 08:00:00.001+0800",
             match serde_json::Value::from(Timestamp::new(1, TimeUnit::Millisecond)) {
                 Value::String(s) => s,
                 _ => unreachable!(),
@@ -676,7 +682,7 @@ mod tests {
         );
 
         assert_eq!(
-            "1970-01-01 00:00:00.000001+0000",
+            "1970-01-01 08:00:00.000001+0800",
             match serde_json::Value::from(Timestamp::new(1, TimeUnit::Microsecond)) {
                 Value::String(s) => s,
                 _ => unreachable!(),
@@ -684,7 +690,7 @@ mod tests {
         );
 
         assert_eq!(
-            "1970-01-01 00:00:00.000000001+0000",
+            "1970-01-01 08:00:00.000000001+0800",
             match serde_json::Value::from(Timestamp::new(1, TimeUnit::Nanosecond)) {
                 Value::String(s) => s,
                 _ => unreachable!(),
@@ -869,4 +875,18 @@ mod tests {
         assert_eq!(1, res.value);
         assert_eq!(TimeUnit::Second, res.unit);
     }
+
+    #[test]
+    fn test_parse_in_time_zone() {
+        std::env::set_var("TZ", "Asia/Shanghai");
+        assert_eq!(
+            Timestamp::new(0, TimeUnit::Nanosecond),
+            Timestamp::from_str("1970-01-01 08:00:00.000").unwrap()
+        );
+
+        assert_eq!(
+            Timestamp::new(0, TimeUnit::Second),
+            Timestamp::from_str("1970-01-01 08:00:00").unwrap()
+        );
+    }
 }

src/datatypes/src/timestamp.rs

@@ -125,11 +125,12 @@ mod tests {
 
     #[test]
     fn test_to_serde_json_value() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         let ts = TimestampSecond::new(123);
         let val = serde_json::Value::from(ts);
         match val {
             serde_json::Value::String(s) => {
-                assert_eq!("1970-01-01 00:02:03+0000", s);
+                assert_eq!("1970-01-01 08:02:03+0800", s);
             }
             _ => unreachable!(),
         }

src/datatypes/src/value.rs

@@ -1345,6 +1345,7 @@ mod tests {
 
     #[test]
     fn test_display() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         assert_eq!(Value::Null.to_string(), "Null");
         assert_eq!(Value::UInt8(8).to_string(), "8");
         assert_eq!(Value::UInt16(16).to_string(), "16");
@@ -1366,11 +1367,11 @@ mod tests {
         assert_eq!(Value::Date(Date::new(0)).to_string(), "1970-01-01");
         assert_eq!(
             Value::DateTime(DateTime::new(0)).to_string(),
-            "1970-01-01 00:00:00"
+            "1970-01-01 08:00:00+0800"
         );
         assert_eq!(
             Value::Timestamp(Timestamp::new(1000, TimeUnit::Millisecond)).to_string(),
-            "1970-01-01 00:00:01+0000"
+            "1970-01-01 08:00:01+0800"
         );
         assert_eq!(
             Value::List(ListValue::new(

src/datatypes/src/vectors/datetime.rs

@@ -37,6 +37,7 @@ mod tests {
 
     #[test]
     fn test_datetime_vector() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         let v = DateTimeVector::new(PrimitiveArray::from_slice([1, 2, 3]));
         assert_eq!(ConcreteDataType::datetime_datatype(), v.data_type());
         assert_eq!(3, v.len());
@@ -63,7 +64,7 @@ mod tests {
             unreachable!()
         }
         assert_eq!(
-            "[\"1970-01-01 00:00:01\",\"1970-01-01 00:00:02\",\"1970-01-01 00:00:03\"]",
+            "[\"1970-01-01 08:00:01+0800\",\"1970-01-01 08:00:02+0800\",\"1970-01-01 08:00:03+0800\"]",
             serde_json::to_string(&v.serialize_to_json().unwrap()).unwrap()
         );
     }

src/servers/src/postgres/handler.rs

@@ -16,7 +16,6 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use chrono::LocalResult;
 use common_query::Output;
 use common_recordbatch::error::Result as RecordBatchResult;
 use common_recordbatch::RecordBatch;
@@ -178,7 +177,7 @@ fn encode_value(value: &Value, builder: &mut DataRowEncoder) -> PgWireResult<()>
             }
         }
         Value::Timestamp(v) => {
-            if let LocalResult::Single(datetime) = v.to_chrono_datetime() {
+            if let Some(datetime) = v.to_chrono_datetime() {
                 builder.encode_field(&datetime)
             } else {
                 Err(PgWireError::ApiError(Box::new(Error::Internal {

src/sql/src/statements.rs

@@ -487,15 +487,16 @@ mod tests {
 
     #[test]
     pub fn test_parse_datetime_literal() {
+        std::env::set_var("TZ", "Asia/Shanghai");
         let value = sql_value_to_value(
             "datetime_col",
             &ConcreteDataType::datetime_datatype(),
-            &SqlValue::DoubleQuotedString("2022-02-22 00:01:03".to_string()),
+            &SqlValue::DoubleQuotedString("2022-02-22 00:01:03+0800".to_string()),
         )
         .unwrap();
         assert_eq!(ConcreteDataType::datetime_datatype(), value.data_type());
         if let Value::DateTime(d) = value {
-            assert_eq!("2022-02-22 00:01:03", d.to_string());
+            assert_eq!("2022-02-22 00:01:03+0800", d.to_string());
         } else {
             unreachable!()
         }