Parse @name as a parameter in the Trino dialect

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')",
+        );
+    }
+}

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 <name> should treat subsequent data as raw (ClickHouse).
     /// When true, after tokenizing `INSERT ... FORMAT <non-VALUES-name>`, 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() == '#')