Enable session time zone override for MySQL (#387)

Author: sgrebnov

src/duckdb/creator.rs

@@ -297,6 +297,7 @@ impl TableManager {
 
     /// Inserts data from this table into the target table.
     #[tracing::instrument(level = "debug", skip_all)]
+    #[allow(dead_code)]
     pub(crate) fn insert_into(
         &self,
         table: &TableManager,

src/sql/arrow_sql_gen/mysql.rs

@@ -15,7 +15,6 @@ use chrono::{NaiveDate, NaiveTime, Timelike};
 use mysql_async::{consts::ColumnFlags, consts::ColumnType, FromValueError, Row, Value};
 use snafu::{ResultExt, Snafu};
 use std::{convert, sync::Arc};
-use time::PrimitiveDateTime;
 
 #[derive(Debug, Snafu)]
 pub enum Error {
@@ -527,7 +526,7 @@ pub fn rows_to_arrow(rows: &[Row], projected_schema: &Option<SchemaRef>) -> Resu
                         .fail();
                     };
                     let v = match handle_null_error(
-                        row.get_opt::<PrimitiveDateTime, usize>(i).transpose(),
+                        row.get_opt::<chrono::NaiveDateTime, usize>(i).transpose(),
                     ) {
                         Ok(v) => v,
                         Err(err) => {
@@ -546,10 +545,7 @@ pub fn rows_to_arrow(rows: &[Row], projected_schema: &Option<SchemaRef>) -> Resu
 
                     match v {
                         Some(v) => {
-                            #[allow(clippy::cast_possible_truncation)]
-                            let timestamp_micros =
-                                (v.assume_utc().unix_timestamp_nanos() / 1_000) as i64;
-                            builder.append_value(timestamp_micros);
+                            builder.append_value(v.and_utc().timestamp_micros());
                         }
                         None => builder.append_null(),
                     }

src/sql/db_connection_pool/mysqlpool.rs

@@ -57,12 +57,16 @@ pub struct MySQLConnectionPool {
     join_push_down: JoinPushDown,
 }
 
-const SETUP_QUERIES: [&str; 4] = [
-    "SET time_zone = '+00:00'",
-    "SET character_set_results = 'utf8mb4'",
-    "SET character_set_client = 'utf8mb4'",
-    "SET character_set_connection = 'utf8mb4'",
-];
+/// Returns the setup queries for the MySQL connection, optionally overriding default time zone (UTC).
+fn get_setup_queries(time_zone: Option<&str>) -> Vec<String> {
+    let tz = time_zone.unwrap_or("+00:00");
+    vec![
+        format!("SET time_zone = '{tz}'"),
+        "SET character_set_results = 'utf8mb4'".to_string(),
+        "SET character_set_client = 'utf8mb4'".to_string(),
+        "SET character_set_connection = 'utf8mb4'".to_string(),
+    ]
+}
 
 impl MySQLConnectionPool {
     /// Creates a new instance of `MySQLConnectionPool`.
@@ -80,6 +84,7 @@ impl MySQLConnectionPool {
     ///   * `sslrootcert` - The path to the root certificate to use when connecting to the MySQL database.
     ///   * `pool_min` - The minimum number of connections to keep open in the pool, lazily created when requested.
     ///   * `pool_max` - The maximum number of connections to allow in the pool.
+    ///   * `time_zone` - The time zone to use for the MySQL connection (e.g., "+2:00", "UTC", etc.). Default is "+00:00" (UTC).
     ///
     /// # Errors
     ///
@@ -167,7 +172,9 @@ impl MySQLConnectionPool {
 
         connection_string = connection_string.ssl_opts(ssl_opts);
 
-        connection_string = connection_string.setup(SETUP_QUERIES.to_vec());
+        connection_string = connection_string.setup(get_setup_queries(
+            params.get("time_zone").map(SecretBox::expose_secret),
+        ));
 
         let opts = mysql_async::Opts::from(connection_string);
 

tests/docker/mod.rs

@@ -134,7 +134,7 @@ impl<'a> ContainerRunner<'a> {
                 format!("{container_port}/tcp"),
                 Some(vec![PortBinding {
                     host_ip: Some("127.0.0.1".to_string()),
-                    host_port: Some(format!("{host_port}/tcp")),
+                    host_port: Some(format!("{host_port}")),
                 }]),
             );
         }

tests/mysql/common.rs

@@ -12,7 +12,7 @@ use crate::{
 const MYSQL_ROOT_PASSWORD: &str = "integration-test-pw";
 const MYSQL_DOCKER_CONTAINER: &str = "runtime-integration-test-mysql";
 
-fn get_mysql_params(port: usize) -> HashMap<String, SecretString> {
+fn get_mysql_params(port: usize, time_zone: Option<&str>) -> HashMap<String, SecretString> {
     let mut params = HashMap::new();
     params.insert(
         "mysql_host".to_string(),
@@ -46,6 +46,9 @@ fn get_mysql_params(port: usize) -> HashMap<String, SecretString> {
         "mysql_pool_max".to_string(),
         SecretString::from("10".to_string()),
     );
+    if let Some(tz) = time_zone {
+        params.insert("mysql_time_zone".to_string(), SecretString::from(tz));
+    }
     params
 }
 
@@ -87,8 +90,9 @@ pub async fn start_mysql_docker_container(port: usize) -> Result<RunningContaine
 #[instrument]
 pub(super) async fn get_mysql_connection_pool(
     port: usize,
+    time_zone: Option<&str>,
 ) -> Result<MySQLConnectionPool, anyhow::Error> {
-    let mysql_pool = MySQLConnectionPool::new(get_mysql_params(port))
+    let mysql_pool = MySQLConnectionPool::new(get_mysql_params(port, time_zone))
         .await
         .expect("Failed to create MySQL Connection Pool");
 

tests/mysql/mod.rs

@@ -109,6 +109,75 @@ VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
+    )
+    .await;
+}
+
+/// Tests the MySQL TIMESTAMP with time zone override.
+/// The test verifies that the TIMESTAMP type correctly adjusts to the specified time zone when retrieved from the database.
+/// `TIMESTAMP` columns should be automatically converted to the specified time zone,
+/// while `DATETIME` columns should remain unchanged.
+async fn test_mysql_timestamp_tz_override(port: usize) {
+    let create_table_stmt = "
+        CREATE TABLE timestamp_tz_table (
+            ts TIMESTAMP,
+            dt DATETIME
+        );
+    ";
+    // values will be inserted in UTC
+    let insert_table_stmt = "
+        INSERT INTO timestamp_tz_table (ts, dt)
+        VALUES
+            ('2024-09-12 10:00:00', '2024-09-12 10:00:00');
+    ";
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("ts", DataType::Timestamp(TimeUnit::Microsecond, None), true),
+        Field::new("dt", DataType::Timestamp(TimeUnit::Microsecond, None), true),
+    ]));
+
+    // Both columns should remain unchanged as target time zone is UTC (same as insert time zone)
+    let expected_utc = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![
+            Arc::new(TimestampMicrosecondArray::from(vec![1_726_135_200_000_000])),
+            Arc::new(TimestampMicrosecondArray::from(vec![1_726_135_200_000_000])),
+        ],
+    )
+    .expect("Failed to created arrow record batch");
+
+    arrow_mysql_one_way(
+        port,
+        "timestamp_tz_table",
+        create_table_stmt,
+        insert_table_stmt,
+        expected_utc,
+        Some("UTC"),
+    )
+    .await;
+
+    // "+02:00"
+    let expected_custom_tz = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![
+            // ts: TIMESTAMP column, should be shifted +2 hours (timezone override)
+            Arc::new(TimestampMicrosecondArray::from(vec![
+                1_726_135_200_000_000 + 2 * 60 * 60 * 1_000_000,
+            ])),
+            // dt: DATETIME column, should remain unchanged
+            Arc::new(TimestampMicrosecondArray::from(vec![1_726_135_200_000_000])),
+        ],
+    )
+    .expect("Failed to created arrow record batch");
+
+    arrow_mysql_one_way(
+        port,
+        "timestamp_tz_table",
+        create_table_stmt,
+        insert_table_stmt,
+        expected_custom_tz,
+        Some("+02:00"), // Override time zone to +02:00
     )
     .await;
 }
@@ -197,6 +266,7 @@ VALUES (
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -269,6 +339,7 @@ VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -315,6 +386,7 @@ VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -377,6 +449,7 @@ VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -429,6 +502,7 @@ VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -484,6 +558,7 @@ INSERT INTO high_precision_decimal (decimal_values) VALUES
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -559,6 +634,7 @@ async fn test_mysql_zero_date_type(port: usize) {
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -593,6 +669,7 @@ async fn test_mysql_decimal_types_to_decimal128(port: usize) {
         create_table_stmt,
         insert_table_stmt,
         expected_record,
+        None,
     )
     .await;
 }
@@ -603,13 +680,15 @@ async fn arrow_mysql_one_way(
     create_table_stmt: &str,
     insert_table_stmt: &str,
     expected_record: RecordBatch,
+    test_query_tz: Option<&str>,
 ) -> Vec<RecordBatch> {
     tracing::debug!("Running tests on {table_name}");
 
     let ctx = SessionContext::new();
-    let pool = common::get_mysql_connection_pool(port)
+    // For dataset initialization we always use UTC (default) timezone
+    let pool = common::get_mysql_connection_pool(port, None)
         .await
-        .expect("MySQL connection pool should be created");
+        .expect("MySQL connection pool for test table creation should be created");
 
     let db_conn = pool
         .connect_direct()
@@ -625,6 +704,12 @@ async fn arrow_mysql_one_way(
         .await
         .expect("SQL mode should be adjusted");
 
+    // Drop table if already exists
+    let _ = db_conn
+        .execute(format!("DROP TABLE IF EXISTS {table_name}").as_str(), &[])
+        .await
+        .expect("MySQL table should be dropped if exists");
+
     // Create table and insert data into mysql test_table
     let _ = db_conn
         .execute(create_table_stmt, &[])
@@ -636,6 +721,11 @@ async fn arrow_mysql_one_way(
         .await
         .expect("MySQL table data should be inserted");
 
+    // For the test query, use a new connection pool with optional time zone override
+    let pool = common::get_mysql_connection_pool(port, test_query_tz)
+        .await
+        .expect("MySQL connection pool for test query should be created");
+
     // Register datafusion table, test mysql row -> arrow conversion
     let sqltable_pool: Arc<
         dyn DbConnectionPool<mysql_async::Conn, &'static (dyn ToValue + Sync)>
@@ -685,6 +775,7 @@ async fn test_mysql_arrow_oneway() {
     let mysql_container = start_mysql_container(port).await;
 
     test_mysql_timestamp_types(port).await;
+    test_mysql_timestamp_tz_override(port).await;
     test_mysql_datetime_types(port).await;
     test_mysql_time_types(port).await;
     test_mysql_enum_types(port).await;