timestamp serialization with tz

Signed-off-by: Veeupup <code@tanweime.com>

common/datavalues/src/types/data_type.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::any::Any;
+use chrono_tz::Tz;
 use std::collections::BTreeMap;
 
 use common_arrow::arrow::datatypes::DataType as ArrowType;
@@ -125,6 +126,10 @@ pub trait DataType: std::fmt::Debug + Sync + Send + DynClone {
 
     fn create_mutable(&self, capacity: usize) -> Box<dyn MutableColumn>;
     fn create_serializer(&self) -> TypeSerializerImpl;
+    /// work for timestamp serializer
+    fn create_serializer_with_tz(&self, _tz: Tz) -> TypeSerializerImpl {
+        unimplemented!()
+    }
     fn create_deserializer(&self, capacity: usize) -> TypeDeserializerImpl;
 }
 

common/datavalues/src/types/serializations/timestamp.rs

@@ -37,6 +37,10 @@ impl Default for TimestampSerializer {
 }
 
 impl TimestampSerializer {
+    pub fn new_with_tz(tz: Tz) -> Self {
+        Self {tz}
+    }
+
     pub fn to_timestamp(&self, value: &i64) -> DateTime<Tz> {
         value.to_timestamp(&self.tz)
     }

common/datavalues/src/types/type_timestamp.rs

@@ -150,6 +150,10 @@ impl DataType for TimestampType {
         TimestampSerializer::default().into()
     }
 
+    fn create_serializer_with_tz(&self, tz: Tz) -> TypeSerializerImpl {
+        TimestampSerializer::new_with_tz(tz).into()
+    }
+
     fn create_deserializer(&self, capacity: usize) -> TypeDeserializerImpl {
         let tz = "UTC".parse::<Tz>().unwrap();
         TimestampDeserializer {

query/src/servers/clickhouse/writers/query_writer.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::borrow::Cow;
+use chrono_tz::Tz;
 
 use common_io::prelude::FormatSettings;
 use common_base::ProgressValues;
@@ -143,7 +144,16 @@ pub fn to_clickhouse_block(block: DataBlock, format: &FormatSettings) -> Result<
         let column = block.column(column_index);
         let field = block.schema().field(column_index);
         let name = field.name();
-        let serializer = field.data_type().create_serializer();
+        let serializer = if field.data_type().data_type_id() == TypeID::Timestamp {
+            let tz =
+            String::from_utf8(format.timezone.clone()).map_err(|_| ErrorCode::LogicalError("timezone must be set"))?;
+            let tz = tz.parse::<Tz>().map_err(|_| {
+                ErrorCode::InvalidTimezone("Timezone has been checked and should be valid")
+            })?;
+            field.data_type().create_serializer_with_tz(tz)
+        }else {
+            field.data_type().create_serializer()
+        };
         result.append_column(column::new_column(
             name,
             serializer.serialize_clickhouse_format(&column.convert_full_column(), format)?,

query/src/servers/http/v1/http_query_handlers.rs

@@ -202,7 +202,7 @@ pub(crate) async fn query_handler(
         .try_create_query(&query_id, ctx, req)
         .await;
 
-    // TODO(veeupup): get query_ctx's format_settings here
+    // TODO(veeupup): get global query_ctx's format_settings, because we cann't set session settings now
     let format = FormatSettings::default();
     match query {
         Ok(query) => {

query/src/servers/mysql/writers/query_result_writer.rs

@@ -111,7 +111,11 @@ impl<'a, W: std::io::Write> DFQueryResultWriter<'a, W> {
         }
 
         let block = blocks[0].clone();
-        let tz: Tz = "UTC".parse().unwrap();
+        let tz =
+            String::from_utf8(format.timezone.clone()).map_err(|_| ErrorCode::LogicalError("timezone must be set"))?;
+        let tz = tz.parse::<Tz>().map_err(|_| {
+            ErrorCode::InvalidTimezone("Timezone has been checked and should be valid")
+        })?;
         match convert_schema(block.schema()) {
             Err(error) => Self::err(&error, dataset_writer),
             Ok(columns) => {

query/src/sessions/query_ctx_shared.rs

@@ -263,6 +263,7 @@ impl QueryContextShared {
             format.field_delimiter = settings.get_field_delimiter()?;
             format.empty_as_default = settings.get_empty_as_default()? > 0;
             format.skip_header = settings.get_skip_header()? > 0;
+            format.timezone = settings.get_timezone()?;
         }
         Ok(format)
     }