diff --git a/crates/polyglot-sql/src/dialects/trino.rs b/crates/polyglot-sql/src/dialects/trino.rs index 63d22e77..79ffa195 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 } @@ -511,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 24a0ae9f..6b8a08b9 100644 --- a/crates/polyglot-sql/src/tokens.rs +++ b/crates/polyglot-sql/src/tokens.rs @@ -1535,6 +1535,11 @@ 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 }`, 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 /// the next blank line or end of input is consumed as a raw data token. @@ -1582,6 +1587,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 +1851,22 @@ 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 == '_' || c == '$' || c == '#' || c == '@' { + i += 1; + } else { + return c == '('; + } + } + false + } + #[inline] fn peek_next(&self) -> char { if self.current + 1 >= self.size { @@ -2415,7 +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 == '@') + 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() == '#')