From 58eb42d43043b925a98337c87dd9881dbdc6e88a Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Tue, 16 Jun 2026 23:40:39 -0400 Subject: [PATCH 1/6] cover named IANA timezone in TIMESTAMPTZ wrapper pushdown --- .../cubesql/src/compile/test/test_wrapper.rs | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index be61997e50c7a..eac46d6e0fd5b 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -532,6 +532,96 @@ GROUP BY )); } +/// Using TIMESTAMP WITH TIME ZONE with an IANA timezone in wrapper should render proper timestamptz in SQL +#[tokio::test] +async fn test_wrapper_timestamptz_named_timezone() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + // language=PostgreSQL + r#" +SELECT + customer_gender +FROM KibanaSampleDataEcommerce +WHERE + order_date >= TIMESTAMP WITH TIME ZONE '2026-06-14T00:00:00 America/Los_Angeles' + AND +-- This filter should trigger pushdown + LOWER(customer_gender) = 'male' +GROUP BY + 1 +; + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let wrapped_sql = &query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .wrapped_sql + .sql; + assert!(wrapped_sql.contains( + "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T00:00:00.000-07:00'" + )); + assert!(!wrapped_sql.contains("America/Los_Angeles")); +} + +/// A named timezone timestamp without a seconds component should still render proper timestamptz in SQL +#[tokio::test] +async fn test_wrapper_timestamptz_named_timezone_no_seconds() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + // language=PostgreSQL + r#" +SELECT + customer_gender +FROM KibanaSampleDataEcommerce +WHERE + order_date >= TIMESTAMP WITH TIME ZONE '2026-06-14 00:00 America/Los_Angeles' + AND +-- This filter should trigger pushdown + LOWER(customer_gender) = 'male' +GROUP BY + 1 +; + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let wrapped_sql = &query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .wrapped_sql + .sql; + assert!(wrapped_sql.contains( + "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T00:00:00.000-07:00'" + )); + assert!(!wrapped_sql.contains("America/Los_Angeles")); +} + // TODO add more time zones // TODO add more TS syntax variants // TODO add TIMESTAMPTZ variant From 3f95d3b0e47e31103b20fac22a1a805c010de5e1 Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Tue, 16 Jun 2026 23:57:12 -0400 Subject: [PATCH 2/6] render named IANA timezone TIMESTAMPTZ literals in pushdown SQL --- .../cubesql/src/compile/engine/df/wrapper.rs | 60 ++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs index cdd5142376396..6fd70482e9db8 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs @@ -21,7 +21,8 @@ use crate::{ }, CubeError, }; -use chrono::{Days, NaiveDate, SecondsFormat, TimeZone, Utc}; +use chrono::{Days, NaiveDate, NaiveDateTime, SecondsFormat, TimeZone, Utc}; +use chrono_tz::Tz; use cubeclient::models::{V1LoadRequestQuery, V1LoadRequestQueryJoinSubquery}; use datafusion::logical_plan::{ExprVisitable, ExpressionVisitor, Recursion}; use datafusion::{ @@ -1826,6 +1827,56 @@ impl WrappedSelectNode { .map_err(|e| DataFusionError::Internal(format!("Can't generate SQL for type: {}", e))) } + fn parse_named_timezone_timestamp(value: &str) -> Option { + let value = value.trim(); + let separator_index = value + .char_indices() + .rev() + .find(|(_, ch)| ch.is_whitespace()) + .map(|(idx, _)| idx)?; + let (timestamp, timezone) = value.split_at(separator_index); + let timezone = timezone.trim().parse::().ok()?; + let timestamp = timestamp.trim(); + let datetime = [ + "%Y-%m-%dT%H:%M:%S%.f", + "%Y-%m-%d %H:%M:%S%.f", + "%Y-%m-%dT%H:%M", + "%Y-%m-%d %H:%M", + ] + .iter() + .find_map(|format| NaiveDateTime::parse_from_str(timestamp, format).ok())?; + let datetime = timezone.from_local_datetime(&datetime).single()?; + + Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) + } + + fn generate_sql_for_named_timezone_timestamp_cast( + sql_generator: Arc, + expr: &Expr, + data_type: &DataType, + ) -> result::Result, DataFusionError> { + if !matches!(data_type, DataType::Timestamp(_, _)) { + return Ok(None); + } + + if let Expr::Literal(ScalarValue::Utf8(Some(value))) = expr { + if let Some(value) = Self::parse_named_timezone_timestamp(value) { + return sql_generator + .get_sql_templates() + .timestamp_literal_expr(value) + .map(Some) + .map_err(|e| { + DataFusionError::Internal(format!( + "Can't generate SQL for timestamp: {}", + e + )) + }); + } + } + + Ok(None) + } + fn generate_typed_null( sql_generator: Arc, data_type: Option, @@ -2127,6 +2178,13 @@ impl WrappedSelectNode { subqueries, ), Expr::Cast { expr, data_type } => { + if let Some(resulting_sql) = Self::generate_sql_for_named_timezone_timestamp_cast( + sql_generator.clone(), + expr.as_ref(), + &data_type, + )? { + return Ok((resulting_sql, sql_query)); + } let (expr, sql_query) = Self::generate_sql_for_expr( sql_query, sql_generator.clone(), From 74f73fc7b97e0c708b44e28cb3f183201312b3bd Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Wed, 17 Jun 2026 00:35:44 -0400 Subject: [PATCH 3/6] preserve timestamp semantics for named time zones --- .../cubesql/src/compile/engine/df/wrapper.rs | 60 +---- rust/cubesql/cubesql/src/compile/router.rs | 5 +- .../cubesql/src/compile/test/test_wrapper.rs | 47 +++- rust/cubesql/cubesql/src/sql/statement.rs | 219 ++++++++++++++++++ 4 files changed, 268 insertions(+), 63 deletions(-) diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs index 6fd70482e9db8..cdd5142376396 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs @@ -21,8 +21,7 @@ use crate::{ }, CubeError, }; -use chrono::{Days, NaiveDate, NaiveDateTime, SecondsFormat, TimeZone, Utc}; -use chrono_tz::Tz; +use chrono::{Days, NaiveDate, SecondsFormat, TimeZone, Utc}; use cubeclient::models::{V1LoadRequestQuery, V1LoadRequestQueryJoinSubquery}; use datafusion::logical_plan::{ExprVisitable, ExpressionVisitor, Recursion}; use datafusion::{ @@ -1827,56 +1826,6 @@ impl WrappedSelectNode { .map_err(|e| DataFusionError::Internal(format!("Can't generate SQL for type: {}", e))) } - fn parse_named_timezone_timestamp(value: &str) -> Option { - let value = value.trim(); - let separator_index = value - .char_indices() - .rev() - .find(|(_, ch)| ch.is_whitespace()) - .map(|(idx, _)| idx)?; - let (timestamp, timezone) = value.split_at(separator_index); - let timezone = timezone.trim().parse::().ok()?; - let timestamp = timestamp.trim(); - let datetime = [ - "%Y-%m-%dT%H:%M:%S%.f", - "%Y-%m-%d %H:%M:%S%.f", - "%Y-%m-%dT%H:%M", - "%Y-%m-%d %H:%M", - ] - .iter() - .find_map(|format| NaiveDateTime::parse_from_str(timestamp, format).ok())?; - let datetime = timezone.from_local_datetime(&datetime).single()?; - - Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) - } - - fn generate_sql_for_named_timezone_timestamp_cast( - sql_generator: Arc, - expr: &Expr, - data_type: &DataType, - ) -> result::Result, DataFusionError> { - if !matches!(data_type, DataType::Timestamp(_, _)) { - return Ok(None); - } - - if let Expr::Literal(ScalarValue::Utf8(Some(value))) = expr { - if let Some(value) = Self::parse_named_timezone_timestamp(value) { - return sql_generator - .get_sql_templates() - .timestamp_literal_expr(value) - .map(Some) - .map_err(|e| { - DataFusionError::Internal(format!( - "Can't generate SQL for timestamp: {}", - e - )) - }); - } - } - - Ok(None) - } - fn generate_typed_null( sql_generator: Arc, data_type: Option, @@ -2178,13 +2127,6 @@ impl WrappedSelectNode { subqueries, ), Expr::Cast { expr, data_type } => { - if let Some(resulting_sql) = Self::generate_sql_for_named_timezone_timestamp_cast( - sql_generator.clone(), - expr.as_ref(), - &data_type, - )? { - return Ok((resulting_sql, sql_query)); - } let (expr, sql_query) = Self::generate_sql_for_expr( sql_query, sql_generator.clone(), diff --git a/rust/cubesql/cubesql/src/compile/router.rs b/rust/cubesql/cubesql/src/compile/router.rs index 76cf2877696ae..b509f57b23f69 100644 --- a/rust/cubesql/cubesql/src/compile/router.rs +++ b/rust/cubesql/cubesql/src/compile/router.rs @@ -17,8 +17,8 @@ use crate::{ dataframe, statement::{ ApproximateCountDistinctVisitor, CastReplacer, RedshiftDatePartReplacer, - SensitiveDataSanitizer, SqlParser062Normalizer, ToTimestampReplacer, - UdfWildcardArgReplacer, + SensitiveDataSanitizer, SqlParser062Normalizer, TimestamptzLiteralReplacer, + ToTimestampReplacer, UdfWildcardArgReplacer, }, ColumnFlags, ColumnType, Session, SessionManager, SessionState, }, @@ -744,6 +744,7 @@ impl QueryRouter { pub fn rewrite_statement(stmt: ast::Statement) -> ast::Statement { let stmt = SqlParser062Normalizer::new().replace(stmt); + let stmt = TimestamptzLiteralReplacer::new().replace(stmt); let stmt = CastReplacer::new().replace(stmt); let stmt = ToTimestampReplacer::new().replace(stmt); let stmt = UdfWildcardArgReplacer::new().replace(stmt); diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index eac46d6e0fd5b..24fc50aac73d4 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -572,7 +572,7 @@ GROUP BY .wrapped_sql .sql; assert!(wrapped_sql.contains( - "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T00:00:00.000-07:00'" + "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T07:00:00.000Z'" )); assert!(!wrapped_sql.contains("America/Los_Angeles")); } @@ -617,11 +617,54 @@ GROUP BY .wrapped_sql .sql; assert!(wrapped_sql.contains( - "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T00:00:00.000-07:00'" + "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T07:00:00.000Z'" )); assert!(!wrapped_sql.contains("America/Los_Angeles")); } +/// Using plain TIMESTAMP with an IANA timezone-like suffix should not render as timestamptz +#[tokio::test] +async fn test_wrapper_timestamp_named_timezone_not_timestamptz() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + // language=PostgreSQL + r#" +SELECT + customer_gender +FROM KibanaSampleDataEcommerce +WHERE + order_date >= TIMESTAMP '2026-06-14 00:00 America/Los_Angeles' + AND +-- This filter should trigger pushdown + LOWER(customer_gender) = 'male' +GROUP BY + 1 +; + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let wrapped_sql = &query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .wrapped_sql + .sql; + assert!(!wrapped_sql.contains("timestamptz '2026-06-14T00:00:00.000-07:00'")); + assert!(wrapped_sql.contains("CAST($1 AS TIMESTAMP)")); +} + // TODO add more time zones // TODO add more TS syntax variants // TODO add TIMESTAMPTZ variant diff --git a/rust/cubesql/cubesql/src/sql/statement.rs b/rust/cubesql/cubesql/src/sql/statement.rs index af3b33164f3c3..cc9f2acebc91d 100644 --- a/rust/cubesql/cubesql/src/sql/statement.rs +++ b/rust/cubesql/cubesql/src/sql/statement.rs @@ -1,3 +1,5 @@ +use chrono::{Duration, LocalResult, NaiveDateTime, Offset, SecondsFormat, TimeZone, Timelike}; +use chrono_tz::Tz; use itertools::Itertools; use log::trace; use pg_srv::{ @@ -856,6 +858,146 @@ impl CastReplacer { } } +#[derive(Debug)] +pub struct TimestamptzLiteralReplacer {} + +impl TimestamptzLiteralReplacer { + pub fn new() -> Self { + Self {} + } + + pub fn replace(mut self, stmt: ast::Statement) -> ast::Statement { + let mut result = stmt; + + self.visit_statement(&mut result).unwrap(); + + result + } + + fn parse_value_to_str(value: &Value) -> Option<&str> { + match value { + Value::SingleQuotedString(str) | Value::DoubleQuotedString(str) => Some(&str), + _ => None, + } + } + + fn named_timezone_timestamp_to_rfc3339(value: &str) -> Option { + let value = value.trim(); + let separator_index = value + .char_indices() + .rev() + .find(|(_, ch)| ch.is_whitespace()) + .map(|(idx, _)| idx)?; + let (timestamp, timezone) = value.split_at(separator_index); + let timezone = timezone.trim().parse::().ok()?; + let timestamp = timestamp.trim(); + let datetime = [ + "%Y-%m-%dT%H:%M:%S%.f", + "%Y-%m-%d %H:%M:%S%.f", + "%Y-%m-%dT%H:%M", + "%Y-%m-%d %H:%M", + ] + .iter() + .find_map(|format| NaiveDateTime::parse_from_str(timestamp, format).ok())?; + + match timezone.from_local_datetime(&datetime) { + LocalResult::Single(datetime) => { + Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) + } + // PostgreSQL resolves ambiguous local times to the standard-time occurrence. For + // fall-back transitions chrono_tz returns the earlier daylight occurrence first. + LocalResult::Ambiguous(_, datetime) => { + Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) + } + LocalResult::None => { + let before = datetime.checked_sub_signed(Duration::days(1))?; + let before = timezone.from_local_datetime(&before).earliest()?; + Some(Self::format_naive_datetime_with_offset( + datetime, + before.offset().fix(), + )) + } + } + } + + fn format_naive_datetime_with_offset( + datetime: NaiveDateTime, + offset: chrono::FixedOffset, + ) -> String { + format!( + "{}.{:03}{}", + datetime.format("%Y-%m-%dT%H:%M:%S"), + datetime.nanosecond() / 1_000_000, + offset + ) + } + + fn is_timestamptz_data_type(data_type: &ast::DataType) -> bool { + match data_type { + ast::DataType::Timestamp( + _, + ast::TimezoneInfo::WithTimeZone | ast::TimezoneInfo::Tz, + ) => true, + ast::DataType::Custom(name, _) => name.to_string().eq_ignore_ascii_case("timestamptz"), + _ => false, + } + } + + /// Rewrites a named-IANA-timezone timestamp value in place to an offset-bearing RFC3339 + /// string. Returns `true` only when the value was a named-timezone literal that was + /// normalized, so callers can leave unrelated timestamptz values (and their type) untouched. + fn normalize_timestamptz_value(value: &mut ast::Value) -> bool { + if let Some(timestamp) = + Self::parse_value_to_str(value).and_then(Self::named_timezone_timestamp_to_rfc3339) + { + *value = Value::SingleQuotedString(timestamp); + true + } else { + false + } + } + + fn normalize_timestamptz_literal(expr: &mut Expr) -> bool { + if let Expr::Value(value) = expr { + Self::normalize_timestamptz_value(&mut value.value) + } else { + false + } + } +} + +impl<'ast> Visitor<'ast, ConnectionError> for TimestamptzLiteralReplacer { + fn transform_expr(&mut self, expr: &mut Expr) -> Result<(), ConnectionError> { + if let Expr::TypedString(typed_string) = expr { + if Self::is_timestamptz_data_type(&typed_string.data_type) + && Self::normalize_timestamptz_value(&mut typed_string.value.value) + { + typed_string.data_type = ast::DataType::Timestamp(None, ast::TimezoneInfo::None); + } + } + + Ok(()) + } + + fn visit_cast(&mut self, expr: &mut Expr) -> Result<(), ConnectionError> { + if let Expr::Cast { + expr: cast_expr, + data_type, + .. + } = expr + { + self.visit_expr(&mut *cast_expr)?; + if Self::is_timestamptz_data_type(data_type) + && Self::normalize_timestamptz_literal(cast_expr) + { + *data_type = ast::DataType::Timestamp(None, ast::TimezoneInfo::None); + } + } + + Ok(()) + } +} + impl<'ast> Visitor<'ast, ConnectionError> for CastReplacer { fn visit_cast(&mut self, expr: &mut Expr) -> Result<(), ConnectionError> { if let Expr::Cast { @@ -1373,6 +1515,37 @@ mod tests { Ok(()) } + fn run_timestamptz_literal_replacer(input: &str, output: &str) -> Result<(), CubeError> { + let stmt = Parser::parse_sql(&PostgreSqlDialect {}, &input) + .unwrap() + .pop() + .expect("must contain at least one statement"); + + let replacer = TimestamptzLiteralReplacer::new(); + let res = replacer.replace(stmt); + + assert_eq!(res.to_string(), output); + + Ok(()) + } + + /// Asserts the replacer leaves the statement untouched. Compares against the canonical + /// round-trip of the same input so we don't have to hand-write the parser's rendering. + fn run_timestamptz_literal_replacer_unchanged(input: &str) -> Result<(), CubeError> { + let stmt = Parser::parse_sql(&PostgreSqlDialect {}, &input) + .unwrap() + .pop() + .expect("must contain at least one statement"); + let expected = stmt.to_string(); + + let replacer = TimestamptzLiteralReplacer::new(); + let res = replacer.replace(stmt); + + assert_eq!(res.to_string(), expected); + + Ok(()) + } + #[test] fn test_cast_replacer() -> Result<(), CubeError> { run_cast_replacer("SELECT 'pg_class'::regclass", "SELECT 1259")?; @@ -1403,6 +1576,52 @@ mod tests { Ok(()) } + #[test] + fn test_timestamptz_literal_replacer() -> Result<(), CubeError> { + run_timestamptz_literal_replacer( + "SELECT TIMESTAMP WITH TIME ZONE '2026-06-14T00:00:00 America/Los_Angeles'", + "SELECT TIMESTAMP '2026-06-14T00:00:00.000-07:00'", + )?; + run_timestamptz_literal_replacer( + "SELECT CAST('2026-06-14 00:00 America/Los_Angeles' AS TIMESTAMPTZ)", + "SELECT CAST('2026-06-14T00:00:00.000-07:00' AS TIMESTAMP)", + )?; + run_timestamptz_literal_replacer( + "SELECT TIMESTAMP '2026-06-14 00:00 America/Los_Angeles'", + "SELECT TIMESTAMP '2026-06-14 00:00 America/Los_Angeles'", + )?; + run_timestamptz_literal_replacer( + "SELECT TIMESTAMP WITH TIME ZONE '2026-11-01 01:30 America/Los_Angeles'", + "SELECT TIMESTAMP '2026-11-01T01:30:00.000-08:00'", + )?; + run_timestamptz_literal_replacer( + "SELECT TIMESTAMP WITH TIME ZONE '2026-03-08 02:30 America/Los_Angeles'", + "SELECT TIMESTAMP '2026-03-08T02:30:00.000-08:00'", + )?; + + Ok(()) + } + + #[test] + fn test_timestamptz_literal_replacer_leaves_non_named_unchanged() -> Result<(), CubeError> { + // Casting a column to timestamptz must keep both the expression and the timezone-aware + // type; only named-IANA-timezone string literals should be rewritten. + run_timestamptz_literal_replacer_unchanged("SELECT CAST(order_date AS TIMESTAMPTZ)")?; + run_timestamptz_literal_replacer_unchanged( + "SELECT CAST(order_date AS TIMESTAMP WITH TIME ZONE)", + )?; + // A timestamptz literal with a numeric offset is already unambiguous and must not be + // downgraded to a plain TIMESTAMP. + run_timestamptz_literal_replacer_unchanged( + "SELECT TIMESTAMP WITH TIME ZONE '2026-06-14 00:00:00+02:00'", + )?; + run_timestamptz_literal_replacer_unchanged( + "SELECT CAST('2026-06-14 00:00:00+02:00' AS TIMESTAMPTZ)", + )?; + + Ok(()) + } + fn run_redshift_date_part_replacer(input: &str, output: &str) -> Result<(), CubeError> { let stmts = Parser::parse_sql(&PostgreSqlDialect {}, &input).unwrap(); From cb16f18e5c2e34dca6ee97bc2861bef354a8321e Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Wed, 17 Jun 2026 10:26:29 -0400 Subject: [PATCH 4/6] gate by template support --- .../cubesql/src/compile/engine/df/wrapper.rs | 40 ++++ .../src/compile/rewrite/rules/wrapper/cast.rs | 98 +++++++++- rust/cubesql/cubesql/src/compile/router.rs | 6 +- .../cubesql/src/compile/test/test_wrapper.rs | 112 ++++++++++- rust/cubesql/cubesql/src/sql/statement.rs | 182 ++++++------------ rust/cubesql/cubesql/src/transport/service.rs | 16 ++ rust/cubesql/cubesql/src/utils/mod.rs | 23 +++ 7 files changed, 334 insertions(+), 143 deletions(-) diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs index cdd5142376396..64e5dd111962c 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs @@ -19,6 +19,7 @@ use crate::{ AliasedColumn, DataSource, LoadRequestMeta, MetaContext, SpanId, SqlGenerator, SqlTemplates, TransportLoadRequestQuery, TransportService, }, + utils::{parse_named_timezone_timestamp, TIMESTAMP_TZ_NAMED_TIMEZONE_CAST_TEMPLATE}, CubeError, }; use chrono::{Days, NaiveDate, SecondsFormat, TimeZone, Utc}; @@ -1816,6 +1817,38 @@ impl WrappedSelectNode { .map_err(|e| DataFusionError::Internal(format!("Can't generate SQL for cast: {}", e))) } + fn generate_sql_for_named_timezone_timestamp_cast( + sql_generator: Arc, + expr: &Expr, + data_type: &DataType, + ) -> result::Result, DataFusionError> { + // DataFusion erases WITH TIME ZONE for these literals. + if !matches!(data_type, DataType::Timestamp(_, _)) { + return Ok(None); + } + + let sql_templates = sql_generator.get_sql_templates(); + if !sql_templates.contains_template(TIMESTAMP_TZ_NAMED_TIMEZONE_CAST_TEMPLATE) { + return Ok(None); + } + + if let Expr::Literal(ScalarValue::Utf8(Some(value))) = expr { + if let Some((timestamp, timezone)) = parse_named_timezone_timestamp(value) { + return sql_templates + .timestamp_tz_named_timezone_cast_expr(timestamp, timezone, value.clone()) + .map(Some) + .map_err(|e| { + DataFusionError::Internal(format!( + "Can't generate SQL for timestamp with named timezone cast: {}", + e + )) + }); + } + } + + Ok(None) + } + fn generate_sql_type( sql_generator: Arc, data_type: DataType, @@ -2127,6 +2160,13 @@ impl WrappedSelectNode { subqueries, ), Expr::Cast { expr, data_type } => { + if let Some(resulting_sql) = Self::generate_sql_for_named_timezone_timestamp_cast( + sql_generator.clone(), + expr.as_ref(), + &data_type, + )? { + return Ok((resulting_sql, sql_query)); + } let (expr, sql_query) = Self::generate_sql_for_expr( sql_query, sql_generator.clone(), diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cast.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cast.rs index 222822db79de5..848c7c6934542 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cast.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cast.rs @@ -1,15 +1,50 @@ -use crate::compile::rewrite::{ - cast_expr, rewrite, rewriter::CubeRewrite, rules::wrapper::WrapperRules, - wrapper_pullup_replacer, wrapper_pushdown_replacer, +use crate::{ + compile::rewrite::{ + cast_expr, rewrite, rewriter::CubeEGraph, rewriter::CubeRewrite, + rules::wrapper::WrapperRules, transforming_rewrite, wrapper_pullup_replacer, + wrapper_pushdown_replacer, wrapper_replacer_context, CastExprDataType, LiteralExprValue, + LogicalPlanLanguage, + }, + transport::DataSource, + utils::{parse_named_timezone_timestamp, TIMESTAMP_TZ_NAMED_TIMEZONE_CAST_TEMPLATE}, }; +use crate::{var, var_iter}; +use datafusion::{arrow::datatypes::DataType, scalar::ScalarValue}; +use egg::{Id, Subst}; impl WrapperRules { pub fn cast_rules(&self, rules: &mut Vec) { rules.extend(vec![ - rewrite( + transforming_rewrite( "wrapper-push-down-cast", - wrapper_pushdown_replacer(cast_expr("?expr", "?data_type"), "?context"), - cast_expr(wrapper_pushdown_replacer("?expr", "?context"), "?data_type"), + wrapper_pushdown_replacer( + cast_expr("?expr", "?data_type"), + wrapper_replacer_context( + "?alias_to_cube", + "?push_to_cube", + "?in_projection", + "?cube_members", + "?grouped_subqueries", + "?ungrouped_scan", + "?input_data_source", + ), + ), + cast_expr( + wrapper_pushdown_replacer( + "?expr", + wrapper_replacer_context( + "?alias_to_cube", + "?push_to_cube", + "?in_projection", + "?cube_members", + "?grouped_subqueries", + "?ungrouped_scan", + "?input_data_source", + ), + ), + "?data_type", + ), + self.transform_cast_pushdown("?input_data_source", "?expr", "?data_type"), ), rewrite( "wrapper-pull-up-cast", @@ -18,4 +53,55 @@ impl WrapperRules { ), ]); } + + fn transform_cast_pushdown( + &self, + input_data_source_var: &str, + expr_var: &str, + data_type_var: &str, + ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { + let input_data_source_var = var!(input_data_source_var); + let expr_var = var!(expr_var); + let data_type_var = var!(data_type_var); + let meta = self.meta_context.clone(); + move |egraph, subst| { + let is_named_timezone_timestamp_cast = + var_iter!(egraph[subst[data_type_var]], CastExprDataType).any(|data_type| { + // DataFusion erases WITH TIME ZONE for these literals. + matches!(data_type, DataType::Timestamp(_, _)) + && Self::expr_is_named_timezone_string_literal(egraph, subst[expr_var]) + }); + + if !is_named_timezone_timestamp_cast { + return true; + } + + let Ok(data_source) = Self::get_data_source(egraph, subst, input_data_source_var) + else { + return false; + }; + + match data_source { + DataSource::Specific(_) => Self::can_rewrite_template( + &data_source, + &meta, + TIMESTAMP_TZ_NAMED_TIMEZONE_CAST_TEMPLATE, + ), + // This template is target-specific. + DataSource::Unrestricted => false, + } + } + } + + fn expr_is_named_timezone_string_literal(egraph: &CubeEGraph, expr_id: Id) -> bool { + egraph[expr_id].nodes.iter().any(|node| { + if let LogicalPlanLanguage::LiteralExpr([value_id]) = node { + var_iter!(egraph[*value_id], LiteralExprValue).any(|literal| { + matches!(literal, ScalarValue::Utf8(Some(value)) if parse_named_timezone_timestamp(value).is_some()) + }) + } else { + false + } + }) + } } diff --git a/rust/cubesql/cubesql/src/compile/router.rs b/rust/cubesql/cubesql/src/compile/router.rs index b509f57b23f69..abba1f5d90bb6 100644 --- a/rust/cubesql/cubesql/src/compile/router.rs +++ b/rust/cubesql/cubesql/src/compile/router.rs @@ -16,8 +16,8 @@ use crate::{ auth_service::SqlAuthServiceAuthenticateRequest, dataframe, statement::{ - ApproximateCountDistinctVisitor, CastReplacer, RedshiftDatePartReplacer, - SensitiveDataSanitizer, SqlParser062Normalizer, TimestamptzLiteralReplacer, + ApproximateCountDistinctVisitor, CastReplacer, PlainTimestampTimezoneSuffixReplacer, + RedshiftDatePartReplacer, SensitiveDataSanitizer, SqlParser062Normalizer, ToTimestampReplacer, UdfWildcardArgReplacer, }, ColumnFlags, ColumnType, Session, SessionManager, SessionState, @@ -744,7 +744,7 @@ impl QueryRouter { pub fn rewrite_statement(stmt: ast::Statement) -> ast::Statement { let stmt = SqlParser062Normalizer::new().replace(stmt); - let stmt = TimestamptzLiteralReplacer::new().replace(stmt); + let stmt = PlainTimestampTimezoneSuffixReplacer::new().replace(stmt); let stmt = CastReplacer::new().replace(stmt); let stmt = ToTimestampReplacer::new().replace(stmt); let stmt = UdfWildcardArgReplacer::new().replace(stmt); diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index 24fc50aac73d4..ee54049d09359 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -24,6 +24,13 @@ use crate::{ CubeError, }; +fn timestamp_tz_named_timezone_cast_template() -> Vec<(String, String)> { + vec![( + "expressions/timestamp_tz_named_timezone_cast".to_string(), + "CONVERT_TIMEZONE('{{ timezone }}', 'UTC', '{{ timestamp }}'::timestamp_ntz)".to_string(), + )] +} + #[tokio::test] async fn test_simple_wrapper() { if !Rewriter::sql_push_down_enabled() { @@ -496,7 +503,7 @@ async fn test_wrapper_timestamptz() { } init_testing_logger(); - let query_plan = convert_select_to_query_plan( + let query_plan = convert_select_to_query_plan_customized( // language=PostgreSQL r#" SELECT @@ -513,6 +520,7 @@ GROUP BY "# .to_string(), DatabaseProtocol::PostgreSQL, + timestamp_tz_named_timezone_cast_template(), ) .await; @@ -540,7 +548,7 @@ async fn test_wrapper_timestamptz_named_timezone() { } init_testing_logger(); - let query_plan = convert_select_to_query_plan( + let query_plan = convert_select_to_query_plan_customized( // language=PostgreSQL r#" SELECT @@ -557,6 +565,7 @@ GROUP BY "# .to_string(), DatabaseProtocol::PostgreSQL, + timestamp_tz_named_timezone_cast_template(), ) .await; @@ -572,9 +581,8 @@ GROUP BY .wrapped_sql .sql; assert!(wrapped_sql.contains( - "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T07:00:00.000Z'" + "${KibanaSampleDataEcommerce.order_date} >= CONVERT_TIMEZONE('America/Los_Angeles', 'UTC', '2026-06-14T00:00:00'::timestamp_ntz)" )); - assert!(!wrapped_sql.contains("America/Los_Angeles")); } /// A named timezone timestamp without a seconds component should still render proper timestamptz in SQL @@ -585,7 +593,7 @@ async fn test_wrapper_timestamptz_named_timezone_no_seconds() { } init_testing_logger(); - let query_plan = convert_select_to_query_plan( + let query_plan = convert_select_to_query_plan_customized( // language=PostgreSQL r#" SELECT @@ -602,6 +610,7 @@ GROUP BY "# .to_string(), DatabaseProtocol::PostgreSQL, + timestamp_tz_named_timezone_cast_template(), ) .await; @@ -617,9 +626,51 @@ GROUP BY .wrapped_sql .sql; assert!(wrapped_sql.contains( - "${KibanaSampleDataEcommerce.order_date} >= timestamptz '2026-06-14T07:00:00.000Z'" + "${KibanaSampleDataEcommerce.order_date} >= CONVERT_TIMEZONE('America/Los_Angeles', 'UTC', '2026-06-14 00:00'::timestamp_ntz)" )); - assert!(!wrapped_sql.contains("America/Los_Angeles")); +} + +#[tokio::test] +async fn test_wrapper_timestamptz_named_timezone_rejects_malformed_timestamp() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan_customized( + // language=PostgreSQL + r#" +SELECT + customer_gender +FROM KibanaSampleDataEcommerce +WHERE + order_date >= TIMESTAMP WITH TIME ZONE '2026-06-14 00:00'' ) OR TRUE -- America/Los_Angeles' + AND +-- This filter should trigger pushdown + LOWER(customer_gender) = 'male' +GROUP BY + 1 +; + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + timestamp_tz_named_timezone_cast_template(), + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let wrapped_sql = &query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .wrapped_sql + .sql; + assert!(!wrapped_sql.contains("CONVERT_TIMEZONE")); + assert!(!wrapped_sql.contains("OR TRUE")); } /// Using plain TIMESTAMP with an IANA timezone-like suffix should not render as timestamptz @@ -630,7 +681,7 @@ async fn test_wrapper_timestamp_named_timezone_not_timestamptz() { } init_testing_logger(); - let query_plan = convert_select_to_query_plan( + let query_plan = convert_select_to_query_plan_customized( // language=PostgreSQL r#" SELECT @@ -647,6 +698,7 @@ GROUP BY "# .to_string(), DatabaseProtocol::PostgreSQL, + timestamp_tz_named_timezone_cast_template(), ) .await; @@ -661,10 +713,52 @@ GROUP BY .find_cube_scan_wrapped_sql() .wrapped_sql .sql; - assert!(!wrapped_sql.contains("timestamptz '2026-06-14T00:00:00.000-07:00'")); + assert!(!wrapped_sql.contains("CONVERT_TIMEZONE")); + assert!(!wrapped_sql.contains("America/Los_Angeles")); assert!(wrapped_sql.contains("CAST($1 AS TIMESTAMP)")); } +#[tokio::test] +async fn test_wrapper_timestamptz_named_timezone_no_template_fallback() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + // language=PostgreSQL + r#" +SELECT + customer_gender +FROM KibanaSampleDataEcommerce +WHERE + order_date >= TIMESTAMP WITH TIME ZONE '2026-06-14T00:00:00 America/Los_Angeles' + AND +-- This filter should trigger pushdown + LOWER(customer_gender) = 'male' +GROUP BY + 1 +; + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); + + let wrapped_sql = &query_plan + .as_logical_plan() + .find_cube_scan_wrapped_sql() + .wrapped_sql + .sql; + assert!(!wrapped_sql.contains("CONVERT_TIMEZONE")); +} + // TODO add more time zones // TODO add more TS syntax variants // TODO add TIMESTAMPTZ variant diff --git a/rust/cubesql/cubesql/src/sql/statement.rs b/rust/cubesql/cubesql/src/sql/statement.rs index cc9f2acebc91d..4083eadc8ab63 100644 --- a/rust/cubesql/cubesql/src/sql/statement.rs +++ b/rust/cubesql/cubesql/src/sql/statement.rs @@ -1,5 +1,3 @@ -use chrono::{Duration, LocalResult, NaiveDateTime, Offset, SecondsFormat, TimeZone, Timelike}; -use chrono_tz::Tz; use itertools::Itertools; use log::trace; use pg_srv::{ @@ -13,7 +11,7 @@ use sqlparser::ast::{ use std::{collections::HashMap, error::Error}; use super::types::ColumnType; -use crate::sql::postgres::ConnectionError; +use crate::{sql::postgres::ConnectionError, utils::parse_named_timezone_timestamp}; #[derive(Debug)] enum PlaceholderType { @@ -859,9 +857,9 @@ impl CastReplacer { } #[derive(Debug)] -pub struct TimestamptzLiteralReplacer {} +pub struct PlainTimestampTimezoneSuffixReplacer {} -impl TimestamptzLiteralReplacer { +impl PlainTimestampTimezoneSuffixReplacer { pub fn new() -> Self { Self {} } @@ -874,105 +872,42 @@ impl TimestamptzLiteralReplacer { result } - fn parse_value_to_str(value: &Value) -> Option<&str> { - match value { - Value::SingleQuotedString(str) | Value::DoubleQuotedString(str) => Some(&str), - _ => None, - } - } - - fn named_timezone_timestamp_to_rfc3339(value: &str) -> Option { - let value = value.trim(); - let separator_index = value - .char_indices() - .rev() - .find(|(_, ch)| ch.is_whitespace()) - .map(|(idx, _)| idx)?; - let (timestamp, timezone) = value.split_at(separator_index); - let timezone = timezone.trim().parse::().ok()?; - let timestamp = timestamp.trim(); - let datetime = [ - "%Y-%m-%dT%H:%M:%S%.f", - "%Y-%m-%d %H:%M:%S%.f", - "%Y-%m-%dT%H:%M", - "%Y-%m-%d %H:%M", - ] - .iter() - .find_map(|format| NaiveDateTime::parse_from_str(timestamp, format).ok())?; - - match timezone.from_local_datetime(&datetime) { - LocalResult::Single(datetime) => { - Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) - } - // PostgreSQL resolves ambiguous local times to the standard-time occurrence. For - // fall-back transitions chrono_tz returns the earlier daylight occurrence first. - LocalResult::Ambiguous(_, datetime) => { - Some(datetime.to_rfc3339_opts(SecondsFormat::Millis, true)) - } - LocalResult::None => { - let before = datetime.checked_sub_signed(Duration::days(1))?; - let before = timezone.from_local_datetime(&before).earliest()?; - Some(Self::format_naive_datetime_with_offset( - datetime, - before.offset().fix(), - )) - } - } - } - - fn format_naive_datetime_with_offset( - datetime: NaiveDateTime, - offset: chrono::FixedOffset, - ) -> String { - format!( - "{}.{:03}{}", - datetime.format("%Y-%m-%dT%H:%M:%S"), - datetime.nanosecond() / 1_000_000, - offset - ) - } - - fn is_timestamptz_data_type(data_type: &ast::DataType) -> bool { - match data_type { + fn is_plain_timestamp_data_type(data_type: &ast::DataType) -> bool { + matches!( + data_type, ast::DataType::Timestamp( _, - ast::TimezoneInfo::WithTimeZone | ast::TimezoneInfo::Tz, - ) => true, - ast::DataType::Custom(name, _) => name.to_string().eq_ignore_ascii_case("timestamptz"), - _ => false, - } + ast::TimezoneInfo::None | ast::TimezoneInfo::WithoutTimeZone + ) + ) } - /// Rewrites a named-IANA-timezone timestamp value in place to an offset-bearing RFC3339 - /// string. Returns `true` only when the value was a named-timezone literal that was - /// normalized, so callers can leave unrelated timestamptz values (and their type) untouched. - fn normalize_timestamptz_value(value: &mut ast::Value) -> bool { - if let Some(timestamp) = - Self::parse_value_to_str(value).and_then(Self::named_timezone_timestamp_to_rfc3339) - { + fn strip_plain_timestamp_timezone_suffix(value: &mut ast::Value) { + // Postgres ignores zones in timestamp without time zone literals. + let timestamp = match value { + Value::SingleQuotedString(str) | Value::DoubleQuotedString(str) => { + parse_named_timezone_timestamp(str).map(|(timestamp, _)| timestamp) + } + _ => None, + }; + + if let Some(timestamp) = timestamp { *value = Value::SingleQuotedString(timestamp); - true - } else { - false } } - fn normalize_timestamptz_literal(expr: &mut Expr) -> bool { + fn strip_plain_timestamp_timezone_literal(expr: &mut Expr) { if let Expr::Value(value) = expr { - Self::normalize_timestamptz_value(&mut value.value) - } else { - false + Self::strip_plain_timestamp_timezone_suffix(&mut value.value); } } } -impl<'ast> Visitor<'ast, ConnectionError> for TimestamptzLiteralReplacer { +impl<'ast> Visitor<'ast, ConnectionError> for PlainTimestampTimezoneSuffixReplacer { fn transform_expr(&mut self, expr: &mut Expr) -> Result<(), ConnectionError> { if let Expr::TypedString(typed_string) = expr { - if Self::is_timestamptz_data_type(&typed_string.data_type) - && Self::normalize_timestamptz_value(&mut typed_string.value.value) - { - typed_string.data_type = ast::DataType::Timestamp(None, ast::TimezoneInfo::None); + if Self::is_plain_timestamp_data_type(&typed_string.data_type) { + Self::strip_plain_timestamp_timezone_suffix(&mut typed_string.value.value); } } @@ -987,10 +922,8 @@ impl<'ast> Visitor<'ast, ConnectionError> for TimestamptzLiteralReplacer { } = expr { self.visit_expr(&mut *cast_expr)?; - if Self::is_timestamptz_data_type(data_type) - && Self::normalize_timestamptz_literal(cast_expr) - { - *data_type = ast::DataType::Timestamp(None, ast::TimezoneInfo::None); + if Self::is_plain_timestamp_data_type(data_type) { + Self::strip_plain_timestamp_timezone_literal(cast_expr); } } @@ -1515,13 +1448,16 @@ mod tests { Ok(()) } - fn run_timestamptz_literal_replacer(input: &str, output: &str) -> Result<(), CubeError> { + fn run_plain_timestamp_timezone_suffix_replacer( + input: &str, + output: &str, + ) -> Result<(), CubeError> { let stmt = Parser::parse_sql(&PostgreSqlDialect {}, &input) .unwrap() .pop() .expect("must contain at least one statement"); - let replacer = TimestamptzLiteralReplacer::new(); + let replacer = PlainTimestampTimezoneSuffixReplacer::new(); let res = replacer.replace(stmt); assert_eq!(res.to_string(), output); @@ -1529,16 +1465,16 @@ mod tests { Ok(()) } - /// Asserts the replacer leaves the statement untouched. Compares against the canonical - /// round-trip of the same input so we don't have to hand-write the parser's rendering. - fn run_timestamptz_literal_replacer_unchanged(input: &str) -> Result<(), CubeError> { + fn run_plain_timestamp_timezone_suffix_replacer_unchanged( + input: &str, + ) -> Result<(), CubeError> { let stmt = Parser::parse_sql(&PostgreSqlDialect {}, &input) .unwrap() .pop() .expect("must contain at least one statement"); let expected = stmt.to_string(); - let replacer = TimestamptzLiteralReplacer::new(); + let replacer = PlainTimestampTimezoneSuffixReplacer::new(); let res = replacer.replace(stmt); assert_eq!(res.to_string(), expected); @@ -1577,45 +1513,41 @@ mod tests { } #[test] - fn test_timestamptz_literal_replacer() -> Result<(), CubeError> { - run_timestamptz_literal_replacer( - "SELECT TIMESTAMP WITH TIME ZONE '2026-06-14T00:00:00 America/Los_Angeles'", - "SELECT TIMESTAMP '2026-06-14T00:00:00.000-07:00'", - )?; - run_timestamptz_literal_replacer( - "SELECT CAST('2026-06-14 00:00 America/Los_Angeles' AS TIMESTAMPTZ)", - "SELECT CAST('2026-06-14T00:00:00.000-07:00' AS TIMESTAMP)", - )?; - run_timestamptz_literal_replacer( - "SELECT TIMESTAMP '2026-06-14 00:00 America/Los_Angeles'", + fn test_plain_timestamp_timezone_suffix_replacer() -> Result<(), CubeError> { + run_plain_timestamp_timezone_suffix_replacer( "SELECT TIMESTAMP '2026-06-14 00:00 America/Los_Angeles'", + "SELECT TIMESTAMP '2026-06-14 00:00'", )?; - run_timestamptz_literal_replacer( - "SELECT TIMESTAMP WITH TIME ZONE '2026-11-01 01:30 America/Los_Angeles'", - "SELECT TIMESTAMP '2026-11-01T01:30:00.000-08:00'", + run_plain_timestamp_timezone_suffix_replacer( + "SELECT CAST('2026-06-14 00:00 America/Los_Angeles' AS TIMESTAMP)", + "SELECT CAST('2026-06-14 00:00' AS TIMESTAMP)", )?; - run_timestamptz_literal_replacer( - "SELECT TIMESTAMP WITH TIME ZONE '2026-03-08 02:30 America/Los_Angeles'", - "SELECT TIMESTAMP '2026-03-08T02:30:00.000-08:00'", + run_plain_timestamp_timezone_suffix_replacer_unchanged( + "SELECT TIMESTAMP 'not a timestamp America/Los_Angeles'", )?; Ok(()) } #[test] - fn test_timestamptz_literal_replacer_leaves_non_named_unchanged() -> Result<(), CubeError> { - // Casting a column to timestamptz must keep both the expression and the timezone-aware - // type; only named-IANA-timezone string literals should be rewritten. - run_timestamptz_literal_replacer_unchanged("SELECT CAST(order_date AS TIMESTAMPTZ)")?; - run_timestamptz_literal_replacer_unchanged( + fn test_plain_timestamp_timezone_suffix_replacer_leaves_timestamptz_unchanged( + ) -> Result<(), CubeError> { + run_plain_timestamp_timezone_suffix_replacer_unchanged( + "SELECT TIMESTAMP WITH TIME ZONE '2026-06-14T00:00:00 America/Los_Angeles'", + )?; + run_plain_timestamp_timezone_suffix_replacer_unchanged( + "SELECT CAST('2026-06-14 00:00 America/Los_Angeles' AS TIMESTAMPTZ)", + )?; + run_plain_timestamp_timezone_suffix_replacer_unchanged( + "SELECT CAST(order_date AS TIMESTAMPTZ)", + )?; + run_plain_timestamp_timezone_suffix_replacer_unchanged( "SELECT CAST(order_date AS TIMESTAMP WITH TIME ZONE)", )?; - // A timestamptz literal with a numeric offset is already unambiguous and must not be - // downgraded to a plain TIMESTAMP. - run_timestamptz_literal_replacer_unchanged( + run_plain_timestamp_timezone_suffix_replacer_unchanged( "SELECT TIMESTAMP WITH TIME ZONE '2026-06-14 00:00:00+02:00'", )?; - run_timestamptz_literal_replacer_unchanged( + run_plain_timestamp_timezone_suffix_replacer_unchanged( "SELECT CAST('2026-06-14 00:00:00+02:00' AS TIMESTAMPTZ)", )?; diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index b2b4b75bf6246..ad15feba79a82 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/rust/cubesql/cubesql/src/transport/service.rs @@ -892,6 +892,22 @@ impl SqlTemplates { self.render_template("expressions/timestamp_literal", context! { value => value }) } + pub fn timestamp_tz_named_timezone_cast_expr( + &self, + timestamp: String, + timezone: String, + value: String, + ) -> Result { + self.render_template( + "expressions/timestamp_tz_named_timezone_cast", + context! { + timestamp => timestamp, + timezone => timezone, + value => value, + }, + ) + } + pub fn like_expr( &self, like_type: LikeType, diff --git a/rust/cubesql/cubesql/src/utils/mod.rs b/rust/cubesql/cubesql/src/utils/mod.rs index 8d3b2cc28246d..11a219fbd222f 100644 --- a/rust/cubesql/cubesql/src/utils/mod.rs +++ b/rust/cubesql/cubesql/src/utils/mod.rs @@ -1,4 +1,6 @@ use crate::compile::rewrite::rewriter::CubeEGraph; +use chrono::NaiveDateTime; +use chrono_tz::Tz; use datafusion::scalar::ScalarValue; use sha2::{Digest, Sha256}; use std::{ @@ -7,6 +9,27 @@ use std::{ hash::{Hash, Hasher}, }; +pub const TIMESTAMP_TZ_NAMED_TIMEZONE_CAST_TEMPLATE: &str = + "expressions/timestamp_tz_named_timezone_cast"; + +pub fn parse_named_timezone_timestamp(value: &str) -> Option<(String, String)> { + let value = value.trim(); + let (timestamp, timezone) = value.rsplit_once(char::is_whitespace)?; + let timestamp = timestamp.trim(); + let timezone = timezone.trim(); + [ + "%Y-%m-%dT%H:%M:%S%.f", + "%Y-%m-%d %H:%M:%S%.f", + "%Y-%m-%dT%H:%M", + "%Y-%m-%d %H:%M", + ] + .iter() + .find_map(|format| NaiveDateTime::parse_from_str(timestamp, format).ok())?; + timezone.parse::().ok()?; + + Some((timestamp.to_string(), timezone.to_string())) +} + pub struct ShaHasher { hasher: Sha256, } From f2c0be577c0b27bead7604d57110ecdd06328355 Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Thu, 18 Jun 2026 00:56:00 -0400 Subject: [PATCH 5/6] wire into snowflakequery --- .../src/adapter/SnowflakeQuery.ts | 1 + .../test/unit/snowflake-query.test.ts | 11 +++++++++++ 2 files changed, 12 insertions(+) create mode 100644 packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts diff --git a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts index c341f025cbb3c..57e0c108235a5 100644 --- a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts +++ b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts @@ -116,6 +116,7 @@ export class SnowflakeQuery extends BaseQuery { templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})'; templates.expressions.interval = 'INTERVAL \'{{ interval }}\''; templates.expressions.timestamp_literal = '\'{{ value }}\'::timestamp_tz'; + templates.expressions.timestamp_tz_named_timezone_cast = 'CONVERT_TIMEZONE(\'{{ timezone }}\', \'UTC\', \'{{ timestamp }}\'::timestamp_ntz)'; templates.expressions.like = '{{ expr }} {% if negated %}NOT {% endif %}LIKE {{ pattern }}{% if default_escape %} ESCAPE \'\\\\\'{% endif %}'; templates.expressions.ilike = '{{ expr }} {% if negated %}NOT {% endif %}ILIKE {{ pattern }}{% if default_escape %} ESCAPE \'\\\\\'{% endif %}'; templates.operators.is_not_distinct_from = 'IS NOT DISTINCT FROM'; diff --git a/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts b/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts new file mode 100644 index 0000000000000..90b765e7c8710 --- /dev/null +++ b/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts @@ -0,0 +1,11 @@ +import { SnowflakeQuery } from '../../src/adapter/SnowflakeQuery'; + +describe('SnowflakeQuery', () => { + it('provides a named timezone timestamptz cast template', () => { + const templates = SnowflakeQuery.prototype.sqlTemplates(); + + expect(templates.expressions.timestamp_tz_named_timezone_cast).toEqual( + 'CONVERT_TIMEZONE(\'{{ timezone }}\', \'UTC\', \'{{ timestamp }}\'::timestamp_ntz)' + ); + }); +}); From 38144c71604ed6198ec780606e0893c81daeefde Mon Sep 17 00:00:00 2001 From: Matthew Orford Date: Thu, 18 Jun 2026 00:58:49 -0400 Subject: [PATCH 6/6] revert useless test --- .../test/unit/snowflake-query.test.ts | 11 ----------- 1 file changed, 11 deletions(-) delete mode 100644 packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts diff --git a/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts b/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts deleted file mode 100644 index 90b765e7c8710..0000000000000 --- a/packages/cubejs-schema-compiler/test/unit/snowflake-query.test.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { SnowflakeQuery } from '../../src/adapter/SnowflakeQuery'; - -describe('SnowflakeQuery', () => { - it('provides a named timezone timestamptz cast template', () => { - const templates = SnowflakeQuery.prototype.sqlTemplates(); - - expect(templates.expressions.timestamp_tz_named_timezone_cast).toEqual( - 'CONVERT_TIMEZONE(\'{{ timezone }}\', \'UTC\', \'{{ timestamp }}\'::timestamp_ntz)' - ); - }); -});