diff --git a/crates/polyglot-sql/Cargo.toml b/crates/polyglot-sql/Cargo.toml index 0bb61700..c3de829b 100644 --- a/crates/polyglot-sql/Cargo.toml +++ b/crates/polyglot-sql/Cargo.toml @@ -61,8 +61,9 @@ dialect-oracle = [] dialect-hive = [] dialect-spark = [] dialect-sqlite = [] -dialect-presto = [] -dialect-trino = [] +# TIMESTAMP(x, zone) checks the zone against `time`'s time-zone list. +dialect-presto = ["time"] +dialect-trino = ["time"] dialect-redshift = [] dialect-databricks = [] dialect-athena = [] diff --git a/crates/polyglot-sql/src/dialects/mod.rs b/crates/polyglot-sql/src/dialects/mod.rs index 19f9c8e7..3e96f577 100644 --- a/crates/polyglot-sql/src/dialects/mod.rs +++ b/crates/polyglot-sql/src/dialects/mod.rs @@ -65,6 +65,11 @@ mod oracle; mod postgres; #[cfg(feature = "dialect-presto")] mod presto; +#[cfg(all( + feature = "transpile", + any(feature = "dialect-presto", feature = "dialect-trino") +))] +mod presto_family; #[cfg(feature = "dialect-redshift")] mod redshift; #[cfg(feature = "dialect-risingwave")] diff --git a/crates/polyglot-sql/src/dialects/presto.rs b/crates/polyglot-sql/src/dialects/presto.rs index 183e1901..ca45bfa8 100644 --- a/crates/polyglot-sql/src/dialects/presto.rs +++ b/crates/polyglot-sql/src/dialects/presto.rs @@ -497,6 +497,8 @@ impl PrestoDialect { fn transform_function(&self, f: Function) -> Result { let name_upper = f.name.to_uppercase(); match name_upper.as_str() { + "TIMESTAMP" => Ok(super::presto_family::timestamp_function_as_cast(f)), + // IFNULL -> COALESCE "IFNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc { original_name: None, diff --git a/crates/polyglot-sql/src/dialects/presto_family.rs b/crates/polyglot-sql/src/dialects/presto_family.rs new file mode 100644 index 00000000..c19637e8 --- /dev/null +++ b/crates/polyglot-sql/src/dialects/presto_family.rs @@ -0,0 +1,45 @@ +//! Rewrites shared by the Presto and Trino dialects, compiled when either is. + +use crate::expressions::{Cast, DataType, Expression, Function, Literal}; + +/// `TIMESTAMP(x)` as `CAST(x AS TIMESTAMP)`, and `TIMESTAMP(x, 'zone')` naming +/// a known time zone as `AT_TIMEZONE(CAST(x AS TIMESTAMP), 'zone')`, as +/// sqlglot's Presto and Trino generators write it: neither engine has a +/// `timestamp()` function. Any other call (a computed or unknown zone, other +/// arities) is returned unchanged, as sqlglot leaves it. +pub(super) fn timestamp_function_as_cast(f: Function) -> Expression { + let is_known_zone = |zone: &Expression| match zone { + Expression::Literal(lit) => match lit.as_ref() { + Literal::String(name) => crate::time::is_valid_timezone(name), + _ => false, + }, + _ => false, + }; + let cast = |this: Expression| { + Expression::Cast(Box::new(Cast { + this, + to: DataType::Timestamp { + precision: None, + timezone: false, + }, + trailing_comments: Vec::new(), + double_colon_syntax: false, + format: None, + default: None, + inferred_type: None, + })) + }; + match f.args.as_slice() { + [_] => cast(f.args.into_iter().next().unwrap()), + [_, zone] if is_known_zone(zone) => { + let mut args = f.args.into_iter(); + let this = cast(args.next().unwrap()); + let zone = args.next().unwrap(); + Expression::Function(Box::new(Function::new( + "AT_TIMEZONE".to_string(), + vec![this, zone], + ))) + } + _ => Expression::Function(Box::new(f)), + } +} diff --git a/crates/polyglot-sql/src/dialects/trino.rs b/crates/polyglot-sql/src/dialects/trino.rs index 63d22e77..efe1744a 100644 --- a/crates/polyglot-sql/src/dialects/trino.rs +++ b/crates/polyglot-sql/src/dialects/trino.rs @@ -202,6 +202,8 @@ impl TrinoDialect { fn transform_function(&self, f: Function) -> Result { let name_upper = f.name.to_uppercase(); match name_upper.as_str() { + "TIMESTAMP" => Ok(super::presto_family::timestamp_function_as_cast(f)), + // IFNULL -> COALESCE "IFNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc { original_name: None, diff --git a/crates/polyglot-sql/tests/trino_timestamp_function.rs b/crates/polyglot-sql/tests/trino_timestamp_function.rs new file mode 100644 index 00000000..4421993a --- /dev/null +++ b/crates/polyglot-sql/tests/trino_timestamp_function.rs @@ -0,0 +1,63 @@ +//! Trino and Presto have no `timestamp()` function. sqlglot reads one as a +//! timestamp conversion and writes it back as a cast (`no_timestamp_sql`), so a +//! query using it runs once sqlglot has reprinted it; these pin that polyglot's +//! Trino and Presto reprints match. + +use polyglot_sql::{transpile, DialectType}; + +fn reprint(sql: &str, dialect: DialectType) -> String { + transpile(sql, dialect, dialect) + .unwrap_or_else(|e| panic!("failed to transpile {sql:?}: {e}")) + .remove(0) +} + +const DIALECTS: [DialectType; 2] = [DialectType::Trino, DialectType::Presto]; + +#[test] +fn timestamp_of_one_argument_is_a_cast() { + for dialect in DIALECTS { + assert_eq!( + reprint("SELECT TIMESTAMP('2025-09-16 13:51:30.000')", dialect), + "SELECT CAST('2025-09-16 13:51:30.000' AS TIMESTAMP)", + ); + assert_eq!( + reprint("SELECT * FROM t WHERE created_at = timestamp(d)", dialect), + "SELECT * FROM t WHERE created_at = CAST(d AS TIMESTAMP)", + ); + } +} + +#[test] +fn timestamp_with_a_literal_zone_is_a_cast_at_that_zone() { + for dialect in DIALECTS { + assert_eq!( + reprint("SELECT TIMESTAMP(d, 'America/New_York')", dialect), + "SELECT AT_TIMEZONE(CAST(d AS TIMESTAMP), 'America/New_York')", + ); + } +} + +#[test] +fn timestamp_with_a_computed_or_unknown_zone_is_left_alone() { + for dialect in DIALECTS { + for sql in [ + "SELECT TIMESTAMP(d, z)", + "SELECT TIMESTAMP(d, '12:00:00')", + "SELECT TIMESTAMP(d, 'Not/AZone')", + ] { + assert_eq!(reprint(sql, dialect), sql); + } + } +} + +#[test] +fn timestamp_types_are_unchanged() { + for dialect in DIALECTS { + for sql in [ + "SELECT CAST(d AS TIMESTAMP)", + "SELECT CAST(d AS TIMESTAMP(3) WITH TIME ZONE)", + ] { + assert_eq!(reprint(sql, dialect), sql); + } + } +}