Add support for Snowflake column aliases that use SQL keywords

Author: yoavcloud

src/dialect/mod.rs

@@ -782,6 +782,13 @@ pub trait Dialect: Debug + Any {
     fn supports_insert_set(&self) -> bool {
         false
     }
+
+    /// Returns true if the specified keyword should be parsed as a select item alias.
+    /// When explicit is true, the keyword is preceded by an `AS` word. Parser is provided
+    /// to enable looking ahead if needed.
+    fn is_select_item_alias(&self, explicit: bool, kw: &Keyword, _parser: &mut Parser) -> bool {
+        explicit || !keywords::RESERVED_FOR_COLUMN_ALIAS.contains(kw)
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/dialect/snowflake.rs

@@ -247,6 +247,50 @@ impl Dialect for SnowflakeDialect {
     fn supports_partiql(&self) -> bool {
         true
     }
+
+    fn is_select_item_alias(&self, explicit: bool, kw: &Keyword, parser: &mut Parser) -> bool {
+        explicit
+            || match kw {
+            // The following keywords can be considered an alias as long as 
+            // they are not followed by other tokens that may change their meaning
+            // e.g. `SELECT * EXCEPT (col1) FROM tbl`
+            Keyword::EXCEPT
+            // e.g. `SELECT 1 LIMIT 5`
+            | Keyword::LIMIT
+            // e.g. `SELECT 1 OFFSET 5 ROWS`
+            | Keyword::OFFSET
+            // e.g. `INSERT INTO t SELECT 1 RETURNING *`
+            | Keyword::RETURNING if !matches!(parser.peek_token_ref().token, Token::Comma | Token::EOF) =>
+            {
+                false
+            }
+
+            // `FETCH` can be considered an alias as long as it's not followed by `FIRST`` or `NEXT`
+            // which would give it a different meanins, for example: `SELECT 1 FETCH FIRST 10 ROWS` - not an alias
+            Keyword::FETCH
+                if parser.peek_keyword(Keyword::FIRST) || parser.peek_keyword(Keyword::NEXT) =>
+            {
+                false
+            }
+
+            // Reserved keywords by the Snowflake dialect, which seem to be less strictive 
+            // than what is listed in `keywords::RESERVED_FOR_COLUMN_ALIAS`. The following 
+            // keywords were tested with the this statement: `SELECT 1 <KW>`.
+            Keyword::FROM
+            | Keyword::GROUP
+            | Keyword::HAVING
+            | Keyword::INTERSECT
+            | Keyword::INTO
+            | Keyword::ORDER
+            | Keyword::SELECT
+            | Keyword::UNION
+            | Keyword::WHERE
+            | Keyword::WITH => false,
+
+            // Any other word is considered an alias
+            _ => true,
+        }
+    }
 }
 
 fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {

src/parser/mod.rs

@@ -8738,6 +8738,34 @@ impl<'a> Parser<'a> {
         Ok(IdentWithAlias { ident, alias })
     }
 
+    // Optionally parses an alias for a select list item
+    fn maybe_parse_select_item_alias(&mut self) -> Result<Option<Ident>, ParserError> {
+        let after_as = self.parse_keyword(Keyword::AS);
+        let next_token = self.next_token();
+        match next_token.token {
+            // Dialect-specific behavior for words that may be reserved from parsed
+            // as select item aliases.
+            Token::Word(w)
+                if self
+                    .dialect
+                    .is_select_item_alias(after_as, &w.keyword, self) =>
+            {
+                Ok(Some(w.into_ident(next_token.span)))
+            }
+            // MSSQL supports single-quoted strings as aliases for columns
+            Token::SingleQuotedString(s) => Ok(Some(Ident::with_quote('\'', s))),
+            // Support for MySql dialect double-quoted string, `AS "HOUR"` for example
+            Token::DoubleQuotedString(s) => Ok(Some(Ident::with_quote('\"', s))),
+            _ => {
+                if after_as {
+                    return self.expected("an identifier after AS", next_token);
+                }
+                self.prev_token();
+                Ok(None) // no alias found
+            }
+        }
+    }
+
     /// Parse `AS identifier` (or simply `identifier` if it's not a reserved keyword)
     /// Some examples with aliases: `SELECT 1 foo`, `SELECT COUNT(*) AS cnt`,
     /// `SELECT ... FROM t1 foo, t2 bar`, `SELECT ... FROM (...) AS bar`
@@ -8756,6 +8784,8 @@ impl<'a> Parser<'a> {
             Token::Word(w) if after_as || !reserved_kwds.contains(&w.keyword) => {
                 Ok(Some(w.into_ident(next_token.span)))
             }
+            // Left the next two patterns for backwards-compatibility (despite not breaking
+            // any tests).
             // MSSQL supports single-quoted strings as aliases for columns
             // We accept them as table aliases too, although MSSQL does not.
             //
@@ -12361,7 +12391,7 @@ impl<'a> Parser<'a> {
                 })
             }
             expr => self
-                .parse_optional_alias(keywords::RESERVED_FOR_COLUMN_ALIAS)
+                .maybe_parse_select_item_alias()
                 .map(|alias| match alias {
                     Some(alias) => SelectItem::ExprWithAlias { expr, alias },
                     None => SelectItem::UnnamedExpr(expr),

tests/sqlparser_snowflake.rs

@@ -3022,3 +3022,32 @@ fn parse_ls_and_rm() {
 
     snowflake().verified_stmt(r#"LIST @"STAGE_WITH_QUOTES""#);
 }
+
+#[test]
+fn test_sql_keywords_as_select_item_aliases() {
+    // Some keywords that should be parsed as an alias
+    let unreserved_kws = vec!["CLUSTER", "FETCH", "RETURNING", "LIMIT", "EXCEPT"];
+    for kw in unreserved_kws {
+        snowflake()
+            .one_statement_parses_to(&format!("SELECT 1 {kw}"), &format!("SELECT 1 AS {kw}"));
+    }
+
+    // Some keywords that should not be parsed as an alias
+    let reserved_kws = vec![
+        "FROM",
+        "GROUP",
+        "HAVING",
+        "INTERSECT",
+        "INTO",
+        "ORDER",
+        "SELECT",
+        "UNION",
+        "WHERE",
+        "WITH",
+    ];
+    for kw in reserved_kws {
+        assert!(snowflake()
+            .parse_sql_statements(&format!("SELECT 1 {kw}"))
+            .is_err());
+    }
+}