From 396ea62e3578baf4d3d21d32fc9be0cd7e6ff75d Mon Sep 17 00:00:00 2001
From: Pawel Leszczynski <leszczynski.pawel@gmail.com>
Date: Thu, 9 Jun 2022 10:45:23 +0200
Subject: [PATCH] snowflake column identifier

Signed-off-by: Pawel Leszczynski <leszczynski.pawel@gmail.com>
---
 src/ast/mod.rs               |  7 +++++++
 src/dialect/snowflake.rs     |  3 ++-
 src/parser.rs                | 17 ++++++++++++++++-
 tests/sqlparser_snowflake.rs |  6 ++++++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index e31701470b..5577135f63 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -217,6 +217,12 @@ pub enum Expr {
     Identifier(Ident),
     /// Multi-part identifier, e.g. `table_alias.column` or `schema.table.col`
     CompoundIdentifier(Vec<Ident>),
+    /// Multi-part identifier with a column number, e.g. [alias.]$file_col_num[.element]
+    CompoundIdentifierWithColumnNumber {
+        alias: Ident,
+        col_num: String,
+        element: Ident
+    },
     /// JSON access (postgres)  eg: data->'tags'
     JsonAccess {
         left: Box<Expr>,
@@ -375,6 +381,7 @@ impl fmt::Display for Expr {
                 Ok(())
             }
             Expr::CompoundIdentifier(s) => write!(f, "{}", display_separated(s, ".")),
+            Expr::CompoundIdentifierWithColumnNumber {alias, col_num, element} => write!(f, "{}.{}.{}", alias, col_num, element),
             Expr::IsTrue(ast) => write!(f, "{} IS TRUE", ast),
             Expr::IsFalse(ast) => write!(f, "{} IS FALSE", ast),
             Expr::IsNull(ast) => write!(f, "{} IS NULL", ast),
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
index 93db956923..188c13b8ed 100644
--- a/src/dialect/snowflake.rs
+++ b/src/dialect/snowflake.rs
@@ -18,7 +18,7 @@ pub struct SnowflakeDialect;
 impl Dialect for SnowflakeDialect {
     // see https://docs.snowflake.com/en/sql-reference/identifiers-syntax.html
     fn is_identifier_start(&self, ch: char) -> bool {
-        ('a'..='z').contains(&ch) || ('A'..='Z').contains(&ch) || ch == '_'
+        ('a'..='z').contains(&ch) || ('A'..='Z').contains(&ch) || ch == '_' || ch == '@'
     }
 
     fn is_identifier_part(&self, ch: char) -> bool {
@@ -27,5 +27,6 @@ impl Dialect for SnowflakeDialect {
             || ('0'..='9').contains(&ch)
             || ch == '$'
             || ch == '_'
+            || ch == ':'
     }
 }
diff --git a/src/parser.rs b/src/parser.rs
index 36c2075c0e..3264e8bef4 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -316,6 +316,9 @@ impl<'a> Parser<'a> {
 
                 while self.consume_token(&Token::Period) {
                     match self.next_token() {
+                        Token::Placeholder(_s) => if id_parts.len()==1 {
+                            // do nothing -> wildcard expression does not occur together with id column number placeholder
+                        }
                         Token::Word(w) => id_parts.push(w.to_ident()),
                         Token::Mul => {
                             return Ok(WildcardExpr::QualifiedWildcard(ObjectName(id_parts)));
@@ -444,8 +447,12 @@ impl<'a> Parser<'a> {
                 _ => match self.peek_token() {
                     Token::LParen | Token::Period => {
                         let mut id_parts: Vec<Ident> = vec![w.to_ident()];
+                        let mut col_num:Option<String> = None;
                         while self.consume_token(&Token::Period) {
                             match self.next_token() {
+                                Token::Placeholder(s) => if id_parts.len()==1 {
+                                    col_num = Some(s);
+                                }
                                 Token::Word(w) => id_parts.push(w.to_ident()),
                                 unexpected => {
                                     return self
@@ -458,7 +465,15 @@ impl<'a> Parser<'a> {
                             self.prev_token();
                             self.parse_function(ObjectName(id_parts))
                         } else {
-                            Ok(Expr::CompoundIdentifier(id_parts))
+                            if !col_num.is_none() && id_parts.len()==2 {
+                                Ok(Expr::CompoundIdentifierWithColumnNumber {
+                                    alias: id_parts[0].clone(),
+                                    col_num: col_num.unwrap(),
+                                    element: id_parts[1].clone()
+                                })
+                            } else {
+                                Ok(Expr::CompoundIdentifier(id_parts))
+                            }
                         }
                     }
                     _ => Ok(Expr::Identifier(w.to_ident())),
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
index c08632a158..78c7a82fa4 100644
--- a/tests/sqlparser_snowflake.rs
+++ b/tests/sqlparser_snowflake.rs
@@ -99,6 +99,12 @@ fn test_single_table_in_parenthesis() {
     );
 }
 
+#[test]
+fn test_snowflake_variables() {
+    let sql = "SELECT t.$1.st AS st FROM @schema.general_finished";
+    snowflake().verified_stmt(sql);
+}
+
 #[test]
 fn test_single_table_in_parenthesis_with_alias() {
     snowflake_and_generic().one_statement_parses_to(