From 2d8d8b81fee26648c1f146adaa0d3ae14f9cf8e8 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 00:55:41 +0000 Subject: [PATCH 1/2] Parse @name as a parameter in the Trino dialect sqlmesh macro variables such as @start_dt tokenized as TSQL-style identifiers and parsed as columns. Trino now emits DAt for @name (not @name( calls), so the parser builds Parameter { style: At }. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019snAi4ViWsw4dALcjcASho --- crates/polyglot-sql/src/dialects/trino.rs | 1 + crates/polyglot-sql/src/tokens.rs | 23 +++++++++++- .../polyglot-sql/tests/trino_at_parameter.rs | 36 +++++++++++++++++++ 3 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 crates/polyglot-sql/tests/trino_at_parameter.rs diff --git a/crates/polyglot-sql/src/dialects/trino.rs b/crates/polyglot-sql/src/dialects/trino.rs index 63d22e77..74aa0205 100644 --- a/crates/polyglot-sql/src/dialects/trino.rs +++ b/crates/polyglot-sql/src/dialects/trino.rs @@ -30,6 +30,7 @@ impl DialectImpl for TrinoDialect { // Trino does NOT support QUALIFY - it's a valid identifier // (unlike Snowflake, BigQuery, DuckDB which have QUALIFY clause) config.keywords.remove("QUALIFY"); + config.at_sign_is_parameter = true; config } diff --git a/crates/polyglot-sql/src/tokens.rs b/crates/polyglot-sql/src/tokens.rs index 24a0ae9f..25641152 100644 --- a/crates/polyglot-sql/src/tokens.rs +++ b/crates/polyglot-sql/src/tokens.rs @@ -1535,6 +1535,10 @@ pub struct TokenizerConfig { /// When true, a bare `$` that is not part of a dollar-quoted string or positional /// parameter is treated as an identifier character. pub dollar_sign_is_identifier: bool, + /// Whether `@name` is a parameter (sqlmesh-style macro variable) rather than a + /// TSQL-style `@variable` identifier. When true, `@` is emitted as `DAt` so the + /// parser builds `Parameter { style: At }`. + pub at_sign_is_parameter: bool, /// Whether INSERT ... FORMAT should treat subsequent data as raw (ClickHouse). /// When true, after tokenizing `INSERT ... FORMAT `, all text until /// the next blank line or end of input is consumed as a raw data token. @@ -1582,6 +1586,7 @@ impl Default for TokenizerConfig { string_escapes_allowed_in_raw_strings: true, hash_comments: false, dollar_sign_is_identifier: false, + at_sign_is_parameter: false, insert_format_raw_data: false, numbers_can_be_underscore_separated: false, recover_terminal_backslash_quote: false, @@ -1845,6 +1850,21 @@ impl<'a, C: TokenizerCursor, T: TokenOutput> TokenizerState<'a, C, T> { } } + /// Whether the `@word` starting at `current` is immediately followed by `(`, + /// i.e. a sqlmesh macro function call such as `@each(...)`. + fn at_word_is_call(&self) -> bool { + let mut i = self.current + 1; + while i < self.size { + let c = self.char_at(i); + if c.is_alphanumeric() || c == '_' { + i += 1; + } else { + return c == '('; + } + } + false + } + #[inline] fn peek_next(&self) -> char { if self.current + 1 >= self.size { @@ -2415,7 +2435,8 @@ impl<'a, C: TokenizerCursor, T: TokenOutput> TokenizerState<'a, C, T> { // TSQL: Check for identifiers starting with # (temp tables) or @ (variables) // e.g., #temp, ##global_temp, @variable - if (c == '#' || c == '@') + if (c == '#' + || (c == '@' && !(self.config.at_sign_is_parameter && !self.at_word_is_call()))) && (self.peek_next().is_alphanumeric() || self.peek_next() == '_' || self.peek_next() == '#') diff --git a/crates/polyglot-sql/tests/trino_at_parameter.rs b/crates/polyglot-sql/tests/trino_at_parameter.rs new file mode 100644 index 00000000..355f0aab --- /dev/null +++ b/crates/polyglot-sql/tests/trino_at_parameter.rs @@ -0,0 +1,36 @@ +use polyglot_sql::expressions::ParameterStyle; +use polyglot_sql::{generate, parse, DialectType, Expression}; + +fn roundtrip_to(sql: &str, expected: &str) -> Expression { + let ast = parse(sql, DialectType::Trino).expect("parse"); + assert_eq!( + generate(&ast[0], DialectType::Trino).expect("generate"), + expected + ); + ast.into_iter().next().unwrap() +} + +#[test] +fn trino_at_name_is_an_at_parameter() { + let sql = "SELECT @start_dt FROM t"; + let Expression::Select(select) = roundtrip_to(sql, sql) else { + panic!("expected Select"); + }; + let Expression::Parameter(param) = &select.expressions[0] else { + panic!( + "expected Parameter, got {}", + select.expressions[0].variant_name() + ); + }; + assert_eq!(param.name.as_deref(), Some("start_dt")); + assert_eq!(param.style, ParameterStyle::At); +} + +#[test] +fn trino_at_macro_function_call_round_trips() { + // Function names are uppercased on regeneration, as before. + roundtrip_to( + "SELECT a FROM t WHERE country = @some_macro_func('x')", + "SELECT a FROM t WHERE country = @SOME_MACRO_FUNC('x')", + ); +} From b4ad7c372e3d55d1228b2772058771797bd303b7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 13 Oct 2026 01:09:03 +0000 Subject: [PATCH 2/2] Review fixes: move Trino @name tests into the lib, tighten call detection Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019snAi4ViWsw4dALcjcASho --- crates/polyglot-sql/src/dialects/trino.rs | 39 +++++++++++++++++++ crates/polyglot-sql/src/tokens.rs | 9 +++-- .../polyglot-sql/tests/trino_at_parameter.rs | 36 ----------------- 3 files changed, 44 insertions(+), 40 deletions(-) delete mode 100644 crates/polyglot-sql/tests/trino_at_parameter.rs diff --git a/crates/polyglot-sql/src/dialects/trino.rs b/crates/polyglot-sql/src/dialects/trino.rs index 74aa0205..79ffa195 100644 --- a/crates/polyglot-sql/src/dialects/trino.rs +++ b/crates/polyglot-sql/src/dialects/trino.rs @@ -512,3 +512,42 @@ impl TrinoDialect { Ok(Expression::Cast(Box::new(c))) } } + +#[cfg(test)] +mod tests { + use crate::expressions::{Expression, ParameterStyle}; + use crate::{generate, parse, DialectType}; + + fn roundtrip_to(sql: &str, expected: &str) -> Expression { + let ast = parse(sql, DialectType::Trino).expect("parse"); + assert_eq!( + generate(&ast[0], DialectType::Trino).expect("generate"), + expected + ); + ast.into_iter().next().unwrap() + } + + #[test] + fn at_name_is_an_at_parameter() { + let sql = "SELECT @start_dt FROM t"; + let Expression::Select(select) = roundtrip_to(sql, sql) else { + panic!("expected Select"); + }; + let Expression::Parameter(param) = &select.expressions[0] else { + panic!( + "expected Parameter, got {}", + select.expressions[0].variant_name() + ); + }; + assert_eq!(param.name.as_deref(), Some("start_dt")); + assert_eq!(param.style, ParameterStyle::At); + } + + #[test] + fn at_macro_function_call_round_trips() { + roundtrip_to( + "SELECT a FROM t WHERE country = @some_macro_func('x')", + "SELECT a FROM t WHERE country = @SOME_MACRO_FUNC('x')", + ); + } +} diff --git a/crates/polyglot-sql/src/tokens.rs b/crates/polyglot-sql/src/tokens.rs index 25641152..6b8a08b9 100644 --- a/crates/polyglot-sql/src/tokens.rs +++ b/crates/polyglot-sql/src/tokens.rs @@ -1537,7 +1537,8 @@ pub struct TokenizerConfig { pub dollar_sign_is_identifier: bool, /// Whether `@name` is a parameter (sqlmesh-style macro variable) rather than a /// TSQL-style `@variable` identifier. When true, `@` is emitted as `DAt` so the - /// parser builds `Parameter { style: At }`. + /// parser builds `Parameter { style: At }`, unless the word is immediately + /// followed by `(` (a macro call), which stays an identifier. pub at_sign_is_parameter: bool, /// Whether INSERT ... FORMAT should treat subsequent data as raw (ClickHouse). /// When true, after tokenizing `INSERT ... FORMAT `, all text until @@ -1852,11 +1853,12 @@ impl<'a, C: TokenizerCursor, T: TokenOutput> TokenizerState<'a, C, T> { /// Whether the `@word` starting at `current` is immediately followed by `(`, /// i.e. a sqlmesh macro function call such as `@each(...)`. + // ponytail: whitespace before `(` (`@each (x)`) is not treated as a call. fn at_word_is_call(&self) -> bool { let mut i = self.current + 1; while i < self.size { let c = self.char_at(i); - if c.is_alphanumeric() || c == '_' { + if c.is_alphanumeric() || c == '_' || c == '$' || c == '#' || c == '@' { i += 1; } else { return c == '('; @@ -2435,8 +2437,7 @@ impl<'a, C: TokenizerCursor, T: TokenOutput> TokenizerState<'a, C, T> { // TSQL: Check for identifiers starting with # (temp tables) or @ (variables) // e.g., #temp, ##global_temp, @variable - if (c == '#' - || (c == '@' && !(self.config.at_sign_is_parameter && !self.at_word_is_call()))) + if (c == '#' || (c == '@' && (!self.config.at_sign_is_parameter || self.at_word_is_call()))) && (self.peek_next().is_alphanumeric() || self.peek_next() == '_' || self.peek_next() == '#') diff --git a/crates/polyglot-sql/tests/trino_at_parameter.rs b/crates/polyglot-sql/tests/trino_at_parameter.rs deleted file mode 100644 index 355f0aab..00000000 --- a/crates/polyglot-sql/tests/trino_at_parameter.rs +++ /dev/null @@ -1,36 +0,0 @@ -use polyglot_sql::expressions::ParameterStyle; -use polyglot_sql::{generate, parse, DialectType, Expression}; - -fn roundtrip_to(sql: &str, expected: &str) -> Expression { - let ast = parse(sql, DialectType::Trino).expect("parse"); - assert_eq!( - generate(&ast[0], DialectType::Trino).expect("generate"), - expected - ); - ast.into_iter().next().unwrap() -} - -#[test] -fn trino_at_name_is_an_at_parameter() { - let sql = "SELECT @start_dt FROM t"; - let Expression::Select(select) = roundtrip_to(sql, sql) else { - panic!("expected Select"); - }; - let Expression::Parameter(param) = &select.expressions[0] else { - panic!( - "expected Parameter, got {}", - select.expressions[0].variant_name() - ); - }; - assert_eq!(param.name.as_deref(), Some("start_dt")); - assert_eq!(param.style, ParameterStyle::At); -} - -#[test] -fn trino_at_macro_function_call_round_trips() { - // Function names are uppercased on regeneration, as before. - roundtrip_to( - "SELECT a FROM t WHERE country = @some_macro_func('x')", - "SELECT a FROM t WHERE country = @SOME_MACRO_FUNC('x')", - ); -}