diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index e1b8413a9..c12722f5c 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -306,6 +306,59 @@ fn mysql_setup_queries(url: &str) -> Vec { vec![format!("SET NAMES {charset}")] } +fn should_enable_explicit_timestamp_defaults(sql: &str) -> bool { + if !starts_with_executable_sql_keyword(sql, &["CREATE", "ALTER"]) { + return false; + } + let lower = sql.split_whitespace().collect::>().join(" ").to_ascii_lowercase(); + lower.contains("timestamp") && lower.contains("default null") +} + +fn explicit_timestamp_defaults_sql(enabled: bool) -> &'static str { + if enabled { + "SET SESSION explicit_defaults_for_timestamp = ON" + } else { + "SET SESSION explicit_defaults_for_timestamp = OFF" + } +} + +async fn enable_explicit_timestamp_defaults_for_query(conn: &mut mysql_async::Conn, sql: &str) -> Option { + if !should_enable_explicit_timestamp_defaults(sql) { + return None; + } + + let previous = match conn.query_first::("SELECT @@SESSION.explicit_defaults_for_timestamp").await { + Ok(Some(value)) => value != 0, + Ok(None) => { + log::debug!("Skipping MySQL explicit timestamp defaults compatibility setting: variable was empty"); + return None; + } + Err(err) => { + log::debug!("Skipping MySQL explicit timestamp defaults compatibility setting: {err}"); + return None; + } + }; + + if previous { + return None; + } + + if let Err(err) = conn.query_drop(explicit_timestamp_defaults_sql(true)).await { + log::debug!("Skipping MySQL explicit timestamp defaults compatibility setting: {err}"); + return None; + } + + Some(previous) +} + +async fn restore_explicit_timestamp_defaults_for_query(conn: &mut mysql_async::Conn, previous: Option) { + if let Some(previous) = previous { + if let Err(err) = conn.query_drop(explicit_timestamp_defaults_sql(previous)).await { + log::warn!("Failed to restore MySQL explicit timestamp defaults session setting: {err}"); + } + } +} + fn mysql_connection_charset(url: &str) -> Option<&str> { let (_, query) = url.split_once('?')?; query.split('&').find_map(|segment| { @@ -710,9 +763,18 @@ pub async fn execute_query_with_max_rows( } } else { let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; - let result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; + let previous_explicit_timestamp_defaults = enable_explicit_timestamp_defaults_for_query(&mut conn, sql).await; + let result = match conn.query_iter(sql).await { + Ok(result) => result, + Err(err) => { + restore_explicit_timestamp_defaults_for_query(&mut conn, previous_explicit_timestamp_defaults).await; + return Err(err.to_string()); + } + }; let affected_rows = result.affected_rows(); - result.drop_result().await.map_err(|e| e.to_string())?; + let drop_result = result.drop_result().await; + restore_explicit_timestamp_defaults_for_query(&mut conn, previous_explicit_timestamp_defaults).await; + drop_result.map_err(|e| e.to_string())?; Ok(QueryResult { columns: vec![], @@ -989,6 +1051,28 @@ mod tests { assert!(!requires_text_protocol_query("SELECT * FROM users")); } + #[test] + fn mysql_timestamp_default_null_ddl_enables_explicit_defaults() { + let create_sql = r#" + CREATE TABLE `referral_record` ( + `id` BINARY(16) NOT NULL, + `created_at` TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), + `updated_at` TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6), + `deleted_at` TIMESTAMP(6) DEFAULT NULL, + PRIMARY KEY (`id`) + ) ENGINE = InnoDB + "#; + + assert!(should_enable_explicit_timestamp_defaults(create_sql)); + assert!(should_enable_explicit_timestamp_defaults( + "ALTER TABLE referral_record ADD deleted_at TIMESTAMP DEFAULT NULL" + )); + assert!(!should_enable_explicit_timestamp_defaults("CREATE TABLE t (deleted_at DATETIME(6) DEFAULT NULL)")); + assert!(!should_enable_explicit_timestamp_defaults("SELECT 'TIMESTAMP DEFAULT NULL'")); + assert_eq!(explicit_timestamp_defaults_sql(true), "SET SESSION explicit_defaults_for_timestamp = ON"); + assert_eq!(explicit_timestamp_defaults_sql(false), "SET SESSION explicit_defaults_for_timestamp = OFF"); + } + #[test] fn mysql_tls_session_close_errors_retry_without_ssl() { let error = "MySQL connection failed: error communicating with database: \