Extended the Snowflake COPY INTO data-load select item parser to support nested semi-structured path traversal

Author: romanoff

src/ast/helpers/stmt_data_loading.rs

@@ -78,8 +78,9 @@ pub struct StageLoadSelectItem {
     pub alias: Option<Ident>,
     /// Column number within the staged file (1-based).
     pub file_col_num: i32,
-    /// Optional element identifier following the column reference.
-    pub element: Option<Ident>,
+    /// Optional semi-structured element path following the column reference
+    /// (e.g. `$1:UsageMetrics:hh` produces `["UsageMetrics", "hh"]`).
+    pub element: Option<Vec<Ident>>,
     /// Optional alias for the item (AS clause).
     pub item_as: Option<Ident>,
 }
@@ -116,8 +117,10 @@ impl fmt::Display for StageLoadSelectItem {
             write!(f, "{alias}.")?;
         }
         write!(f, "${}", self.file_col_num)?;
-        if let Some(element) = &self.element {
-            write!(f, ":{element}")?;
+        if let Some(elements) = &self.element {
+            for element in elements {
+                write!(f, ":{element}")?;
+            }
         }
         if let Some(item_as) = &self.item_as {
             write!(f, " AS {item_as}")?;

src/dialect/snowflake.rs

@@ -37,7 +37,7 @@ use crate::ast::{
     StorageSerializationPolicy, TableObject, TagsColumnOption, Value, WrappedCollection,
 };
 use crate::dialect::{Dialect, Precedence};
-use crate::keywords::Keyword;
+use crate::keywords::{Keyword, RESERVED_FOR_COLUMN_ALIAS};
 use crate::parser::{IsOptional, Parser, ParserError};
 use crate::tokenizer::Token;
 use crate::tokenizer::TokenWithSpan;
@@ -1460,7 +1460,7 @@ fn parse_select_item_for_data_load(
 ) -> Result<StageLoadSelectItem, ParserError> {
     let mut alias: Option<Ident> = None;
     let mut file_col_num: i32 = 0;
-    let mut element: Option<Ident> = None;
+    let mut element: Option<Vec<Ident>> = None;
     let mut item_as: Option<Ident> = None;
 
     let next_token = parser.next_token();
@@ -1493,27 +1493,30 @@ fn parse_select_item_for_data_load(
         }?;
     }
 
-    // try extracting optional element
-    match parser.next_token().token {
-        Token::Colon => {
-            // parse element
-            element = Some(Ident::new(match parser.next_token().token {
-                Token::Word(w) => Ok(w.value),
-                _ => parser.expected("file_col_num", parser.peek_token()),
-            }?));
-        }
-        _ => {
-            // element not present move back
-            parser.prev_token();
+    // try extracting optional element path (e.g. :UsageMetrics:hh)
+    let mut elements = Vec::new();
+    while parser.next_token().token == Token::Colon {
+        match parser.next_token().token {
+            Token::Word(w) => elements.push(Ident::new(w.value)),
+            _ => return parser.expected("element name", parser.peek_token()),
         }
     }
+    parser.prev_token();
+    if !elements.is_empty() {
+        element = Some(elements);
+    }
 
-    // as
+    // optional alias: `AS alias` or just `alias` (implicit)
     if parser.parse_keyword(Keyword::AS) {
         item_as = Some(match parser.next_token().token {
             Token::Word(w) => Ok(Ident::new(w.value)),
             _ => parser.expected("column item alias", parser.peek_token()),
         }?);
+    } else if let Token::Word(w) = parser.peek_token().token {
+        if !RESERVED_FOR_COLUMN_ALIAS.contains(&w.keyword) {
+            parser.next_token();
+            item_as = Some(Ident::new(w.value));
+        }
     }
 
     Ok(StageLoadSelectItem {

tests/sqlparser_snowflake.rs

@@ -2363,7 +2363,7 @@ fn test_copy_into_with_transformations() {
                 StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
                     alias: Some(Ident::new("t1")),
                     file_col_num: 1,
-                    element: Some(Ident::new("st")),
+                    element: Some(vec![Ident::new("st")]),
                     item_as: Some(Ident::new("st"))
                 })
             );
@@ -2372,7 +2372,7 @@ fn test_copy_into_with_transformations() {
                 StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
                     alias: None,
                     file_col_num: 1,
-                    element: Some(Ident::new("index")),
+                    element: Some(vec![Ident::new("index")]),
                     item_as: None
                 })
             );
@@ -2634,6 +2634,28 @@ fn test_snowflake_copy_into_stage_name_ends_with_parens() {
     }
 }
 
+#[test]
+fn test_copy_into_with_nested_colon_path() {
+    let sql = "COPY INTO tbl (col) FROM (SELECT $1:a:b AS col FROM @stage)";
+    match snowflake().verified_stmt(sql) {
+        Statement::CopyIntoSnowflake {
+            from_transformations,
+            ..
+        } => {
+            assert_eq!(
+                from_transformations.as_ref().unwrap()[0],
+                StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
+                    alias: None,
+                    file_col_num: 1,
+                    element: Some(vec![Ident::new("a"), Ident::new("b")]),
+                    item_as: Some(Ident::new("col"))
+                })
+            );
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn test_snowflake_trim() {
     let real_sql = r#"SELECT customer_id, TRIM(sub_items.value:item_price_id, '"', "a") AS item_price_id FROM models_staging.subscriptions"#;