Move stage parsing from tokenizer to parser

apache · Oct 26, 2023 · 02fe126 · 02fe126
1 parent 6e0588c
commit 02fe126
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 18 deletions.
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
@@ -146,8 +146,48 @@ pub fn parse_create_stage(
     })
 }
 
+pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserError> {
+    let mut ident = String::new();
+    while let Some(next_token) = parser.next_token_no_skip() {
+        match &next_token.token {
+            Token::Whitespace(_) => break,
+            Token::Period => {
+                parser.prev_token();
+                break;
+            }
+            Token::AtSign => ident.push('@'),
+            Token::Tilde => ident.push('~'),
+            Token::Mod => ident.push('%'),
+            Token::Div => ident.push('/'),
+            Token::Word(w) => ident.push_str(&w.value),
+            _ => return parser.expected("stage name identifier", parser.peek_token()),
+        }
+    }
+    Ok(Ident::new(ident))
+}
+
+pub fn parse_snowflake_stage_name(parser: &mut Parser) -> Result<ObjectName, ParserError> {
+    match parser.next_token().token {
+        Token::AtSign => {
+            parser.prev_token();
+            let mut idents = vec![];
+            loop {
+                idents.push(parse_stage_name_identifier(parser)?);
+                if !parser.consume_token(&Token::Period) {
+                    break;
+                }
+            }
+            Ok(ObjectName(idents))
+        }
+        _ => {
+            parser.prev_token();
+            Ok(parser.parse_object_name()?)
+        }
+    }
+}
+
 pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
-    let into: ObjectName = parser.parse_object_name()?;
+    let into: ObjectName = parse_snowflake_stage_name(parser)?;
     let mut files: Vec<String> = vec![];
     let mut from_transformations: Option<Vec<StageLoadSelectItem>> = None;
     let from_stage_alias;
@@ -163,7 +203,7 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
             from_transformations = parse_select_items_for_data_load(parser)?;
 
             parser.expect_keyword(Keyword::FROM)?;
-            from_stage = parser.parse_object_name()?;
+            from_stage = parse_snowflake_stage_name(parser)?;
             stage_params = parse_stage_params(parser)?;
 
             // as

diff --git a/src/tokenizer.rs b/src/tokenizer.rs
@@ -1004,18 +1004,6 @@ impl<'a> Tokenizer<'a> {
                             }
                         }
                         Some(' ') => Ok(Some(Token::AtSign)),
-                        // Snowflake stage identifier, this should be consumed as multiple dot separated word tokens
-                        Some(_) if dialect_of!(self is SnowflakeDialect) => {
-                            let mut s = "@".to_string();
-                            s.push_str(&peeking_take_while(chars, |ch| {
-                                self.dialect.is_identifier_part(ch)
-                                    || ch == '/'
-                                    || ch == '~'
-                                    || ch == '%'
-                                    || ch == '.'
-                            }));
-                            Ok(Some(Token::make_word(&s, None)))
-                        }
                         Some(sch) if self.dialect.is_identifier_start('@') => {
                             self.tokenize_identifier_or_keyword([ch, *sch], chars)
                         }

diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
@@ -920,7 +920,7 @@ fn test_copy_into_with_transformations() {
         } => {
             assert_eq!(
                 from_stage,
-                ObjectName(vec![Ident::new("@schema.general_finished")])
+                ObjectName(vec![Ident::new("@schema"), Ident::new("general_finished")])
             );
             assert_eq!(
                 from_transformations.as_ref().unwrap()[0],
@@ -1027,9 +1027,15 @@ fn test_snowflake_stage_object_names() {
     ];
     let mut allowed_object_names = vec![
         ObjectName(vec![Ident::new("my_company"), Ident::new("emp_basic")]),
-        ObjectName(vec![Ident::new("@namespace.%table_name")]),
-        ObjectName(vec![Ident::new("@namespace.%table_name/path")]),
-        ObjectName(vec![Ident::new("@namespace.stage_name/path")]),
+        ObjectName(vec![Ident::new("@namespace"), Ident::new("%table_name")]),
+        ObjectName(vec![
+            Ident::new("@namespace"),
+            Ident::new("%table_name/path"),
+        ]),
+        ObjectName(vec![
+            Ident::new("@namespace"),
+            Ident::new("stage_name/path"),
+        ]),
         ObjectName(vec![Ident::new("@~/path")]),
     ];