diff --git a/src/ast/helpers/stmt_data_loading.rs b/src/ast/helpers/stmt_data_loading.rs index dfc1f4b0b..5cbb0d62a 100644 --- a/src/ast/helpers/stmt_data_loading.rs +++ b/src/ast/helpers/stmt_data_loading.rs @@ -21,6 +21,8 @@ #[cfg(not(feature = "std"))] use alloc::string::String; +#[cfg(not(feature = "std"))] +use alloc::vec::Vec; use core::fmt; #[cfg(feature = "serde")] @@ -78,8 +80,9 @@ pub struct StageLoadSelectItem { pub alias: Option, /// Column number within the staged file (1-based). pub file_col_num: i32, - /// Optional element identifier following the column reference. - pub element: Option, + /// Optional semi-structured element path following the column reference + /// (e.g. `$1:UsageMetrics:hh` produces `["UsageMetrics", "hh"]`). + pub element: Option>, /// Optional alias for the item (AS clause). pub item_as: Option, } @@ -116,8 +119,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}")?; diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs index 31a17225f..0ec7b1337 100644 --- a/src/dialect/snowflake.rs +++ b/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 { let mut alias: Option = None; let mut file_col_num: i32 = 0; - let mut element: Option = None; + let mut element: Option> = None; let mut item_as: Option = 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 { diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index 43444016f..ae07ba159 100644 --- a/tests/sqlparser_snowflake.rs +++ b/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"#;