test: add tests for mysql timestamp encoding

GreptimeTeam · Nov 23, 2023 · 92998ee · 92998ee
1 parent b8d546a
commit 92998ee
Show file tree

Hide file tree

Showing 7 changed files with 239 additions and 67 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/time/src/datetime.rs b/src/common/time/src/datetime.rs
@@ -15,10 +15,11 @@
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
 
-use chrono::{LocalResult, NaiveDateTime};
+use chrono::{LocalResult, NaiveDateTime, TimeZone as ChronoTimeZone, Utc};
 use serde::{Deserialize, Serialize};
 
 use crate::error::{Error, InvalidDateStrSnafu, Result};
+use crate::timezone::TimeZone;
 use crate::util::{format_utc_datetime, local_datetime_to_utc};
 use crate::Date;
 
@@ -108,6 +109,15 @@ impl DateTime {
         NaiveDateTime::from_timestamp_millis(self.0)
     }
 
+    pub fn to_chrono_datetime_with_timezone(&self, tz: Option<TimeZone>) -> Option<NaiveDateTime> {
+        let datetime = self.to_chrono_datetime();
+        datetime.map(|v| match tz {
+            Some(TimeZone::Offset(offset)) => offset.from_utc_datetime(&v).naive_local(),
+            Some(TimeZone::Named(tz)) => tz.from_utc_datetime(&v).naive_local(),
+            None => Utc.from_utc_datetime(&v).naive_local(),
+        })
+    }
+
     /// Convert to [common_time::date].
     pub fn to_date(&self) -> Option<Date> {
         self.to_chrono_datetime().map(|d| Date::from(d.date()))

diff --git a/src/common/time/src/timestamp.rs b/src/common/time/src/timestamp.rs
@@ -21,7 +21,7 @@ use std::time::Duration;
 
 use arrow::datatypes::TimeUnit as ArrowTimeUnit;
 use chrono::{
-    DateTime, LocalResult, NaiveDate, NaiveDateTime, NaiveTime, TimeZone as ChronoTimeZone,
+    DateTime, LocalResult, NaiveDate, NaiveDateTime, NaiveTime, TimeZone as ChronoTimeZone, Utc,
 };
 use serde::{Deserialize, Serialize};
 use snafu::{OptionExt, ResultExt};
@@ -252,6 +252,15 @@ impl Timestamp {
         NaiveDateTime::from_timestamp_opt(sec, nsec)
     }
 
+    pub fn to_chrono_datetime_with_timezone(&self, tz: Option<TimeZone>) -> Option<NaiveDateTime> {
+        let datetime = self.to_chrono_datetime();
+        datetime.map(|v| match tz {
+            Some(TimeZone::Offset(offset)) => offset.from_utc_datetime(&v).naive_local(),
+            Some(TimeZone::Named(tz)) => tz.from_utc_datetime(&v).naive_local(),
+            None => Utc.from_utc_datetime(&v).naive_local(),
+        })
+    }
+
     /// Convert timestamp to chrono date.
     pub fn to_chrono_date(&self) -> Option<NaiveDate> {
         self.to_chrono_datetime().map(|ndt| ndt.date())

diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml
@@ -54,7 +54,7 @@ lazy_static.workspace = true
 mime_guess = "2.0"
 once_cell.workspace = true
 openmetrics-parser = "0.4"
-opensrv-mysql = "0.4"
+opensrv-mysql = "0.5"
 opentelemetry-proto.workspace = true
 parking_lot = "0.12"
 pgwire = "0.16"
@@ -103,7 +103,7 @@ catalog = { workspace = true, features = ["testing"] }
 client.workspace = true
 common-base.workspace = true
 common-test-util.workspace = true
-mysql_async = { git = "https://github.com/blackbeam/mysql_async.git", rev = "32c6f2a986789f97108502c2d0c755a089411b66", default-features = false, features = [
+mysql_async = { version = "0.33", default-features = false, features = [
     "default-rustls",
 ] }
 rand.workspace = true

diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs
@@ -192,8 +192,11 @@ impl<'a, W: AsyncWrite + Unpin> MysqlResultWriter<'a, W> {
                     Value::String(v) => row_writer.write_col(v.as_utf8())?,
                     Value::Binary(v) => row_writer.write_col(v.deref())?,
                     Value::Date(v) => row_writer.write_col(v.to_chrono_date())?,
-                    Value::DateTime(v) => row_writer.write_col(v.to_chrono_datetime())?,
-                    Value::Timestamp(v) => row_writer.write_col(v.to_chrono_datetime())?,
+                    // convert datetime and timestamp to timezone of current connection
+                    Value::DateTime(v) => row_writer
+                        .write_col(v.to_chrono_datetime_with_timezone(query_context.time_zone()))?,
+                    Value::Timestamp(v) => row_writer
+                        .write_col(v.to_chrono_datetime_with_timezone(query_context.time_zone()))?,
                     Value::Interval(v) => row_writer.write_col(v.to_iso8601_string())?,
                     Value::Duration(v) => row_writer.write_col(v.to_std_duration())?,
                     Value::List(_) => {

diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml
@@ -36,6 +36,9 @@ frontend = { workspace = true, features = ["testing"] }
 futures.workspace = true
 meta-client.workspace = true
 meta-srv = { workspace = true, features = ["mock"] }
+mysql_async = { version = "0.33", default-features = false, features = [
+    "default-rustls",
+] }
 object-store.workspace = true
 once_cell.workspace = true
 operator.workspace = true
@@ -59,6 +62,7 @@ sqlx = { version = "0.6", features = [
 substrait.workspace = true
 table.workspace = true
 tempfile.workspace = true
+time = "0.3"
 tokio.workspace = true
 tonic.workspace = true
 tower = "0.4"