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

@@ -1470,8 +1470,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 item_as: Option<Ident> = None;
+    let mut element: Option<Vec<Ident>> = None;
 
     let next_token = parser.next_token();
     match next_token.token {
@@ -1503,29 +1502,21 @@ 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_ref("file_col_num", parser.peek_token_ref()),
-            }?));
-        }
-        _ => {
-            // 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_ref("element name", parser.peek_token_ref()),
         }
     }
-
-    // as
-    if parser.parse_keyword(Keyword::AS) {
-        item_as = Some(match parser.next_token().token {
-            Token::Word(w) => Ok(Ident::new(w.value)),
-            _ => parser.expected_ref("column item alias", parser.peek_token_ref()),
-        }?);
+    parser.prev_token();
+    if !elements.is_empty() {
+        element = Some(elements);
     }
 
+    let item_as = parser.maybe_parse_select_item_alias()?;
+
     Ok(StageLoadSelectItem {
         alias,
         file_col_num,

src/parser/mod.rs

@@ -12314,7 +12314,7 @@ impl<'a> Parser<'a> {
     }
 
     /// Optionally parses an alias for a select list item
-    fn maybe_parse_select_item_alias(&mut self) -> Result<Option<Ident>, ParserError> {
+    pub fn maybe_parse_select_item_alias(&mut self) -> Result<Option<Ident>, ParserError> {
         fn validator(explicit: bool, kw: &Keyword, parser: &mut Parser) -> bool {
             parser.dialect.is_select_item_alias(explicit, kw, parser)
         }

tests/sqlparser_snowflake.rs

@@ -2369,7 +2369,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"))
                 })
             );
@@ -2378,7 +2378,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
                 })
             );
@@ -2640,6 +2640,70 @@ fn test_snowflake_copy_into_stage_name_ends_with_parens() {
     }
 }
 
+#[test]
+fn test_copy_into_with_nested_colon_path() {
+    // Nested colon path with explicit AS alias
+    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!(),
+    }
+
+    // Nested colon path with implicit alias (no AS keyword)
+    let sql = "COPY INTO tbl (col) FROM (SELECT $1:a:b col FROM @stage)";
+    let stmts = snowflake().parse_sql_statements(sql).unwrap();
+    match &stmts[0] {
+        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!(),
+    }
+
+    // Nested colon path with no alias
+    let sql = "COPY INTO tbl FROM (SELECT $1:a:b 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: None
+                })
+            );
+        }
+        _ => 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"#;