Skip to content

Commit 26574b4

Browse files
committed
Extended the Snowflake COPY INTO data-load select item parser to support nested semi-structured path traversal
1 parent 40350e3 commit 26574b4

File tree

3 files changed

+49
-21
lines changed

3 files changed

+49
-21
lines changed

src/ast/helpers/stmt_data_loading.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,9 @@ pub struct StageLoadSelectItem {
7878
pub alias: Option<Ident>,
7979
/// Column number within the staged file (1-based).
8080
pub file_col_num: i32,
81-
/// Optional element identifier following the column reference.
82-
pub element: Option<Ident>,
81+
/// Optional semi-structured element path following the column reference
82+
/// (e.g. `$1:UsageMetrics:hh` produces `["UsageMetrics", "hh"]`).
83+
pub element: Option<Vec<Ident>>,
8384
/// Optional alias for the item (AS clause).
8485
pub item_as: Option<Ident>,
8586
}
@@ -116,8 +117,10 @@ impl fmt::Display for StageLoadSelectItem {
116117
write!(f, "{alias}.")?;
117118
}
118119
write!(f, "${}", self.file_col_num)?;
119-
if let Some(element) = &self.element {
120-
write!(f, ":{element}")?;
120+
if let Some(elements) = &self.element {
121+
for element in elements {
122+
write!(f, ":{element}")?;
123+
}
121124
}
122125
if let Some(item_as) = &self.item_as {
123126
write!(f, " AS {item_as}")?;

src/dialect/snowflake.rs

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ use crate::ast::{
3737
StorageSerializationPolicy, TableObject, TagsColumnOption, Value, WrappedCollection,
3838
};
3939
use crate::dialect::{Dialect, Precedence};
40-
use crate::keywords::Keyword;
40+
use crate::keywords::{Keyword, RESERVED_FOR_COLUMN_ALIAS};
4141
use crate::parser::{IsOptional, Parser, ParserError};
4242
use crate::tokenizer::Token;
4343
use crate::tokenizer::TokenWithSpan;
@@ -1460,7 +1460,7 @@ fn parse_select_item_for_data_load(
14601460
) -> Result<StageLoadSelectItem, ParserError> {
14611461
let mut alias: Option<Ident> = None;
14621462
let mut file_col_num: i32 = 0;
1463-
let mut element: Option<Ident> = None;
1463+
let mut element: Option<Vec<Ident>> = None;
14641464
let mut item_as: Option<Ident> = None;
14651465

14661466
let next_token = parser.next_token();
@@ -1493,27 +1493,30 @@ fn parse_select_item_for_data_load(
14931493
}?;
14941494
}
14951495

1496-
// try extracting optional element
1497-
match parser.next_token().token {
1498-
Token::Colon => {
1499-
// parse element
1500-
element = Some(Ident::new(match parser.next_token().token {
1501-
Token::Word(w) => Ok(w.value),
1502-
_ => parser.expected("file_col_num", parser.peek_token()),
1503-
}?));
1504-
}
1505-
_ => {
1506-
// element not present move back
1507-
parser.prev_token();
1496+
// try extracting optional element path (e.g. :UsageMetrics:hh)
1497+
let mut elements = Vec::new();
1498+
while parser.next_token().token == Token::Colon {
1499+
match parser.next_token().token {
1500+
Token::Word(w) => elements.push(Ident::new(w.value)),
1501+
_ => return parser.expected("element name", parser.peek_token()),
15081502
}
15091503
}
1504+
parser.prev_token();
1505+
if !elements.is_empty() {
1506+
element = Some(elements);
1507+
}
15101508

1511-
// as
1509+
// optional alias: `AS alias` or just `alias` (implicit)
15121510
if parser.parse_keyword(Keyword::AS) {
15131511
item_as = Some(match parser.next_token().token {
15141512
Token::Word(w) => Ok(Ident::new(w.value)),
15151513
_ => parser.expected("column item alias", parser.peek_token()),
15161514
}?);
1515+
} else if let Token::Word(w) = parser.peek_token().token {
1516+
if !RESERVED_FOR_COLUMN_ALIAS.contains(&w.keyword) {
1517+
parser.next_token();
1518+
item_as = Some(Ident::new(w.value));
1519+
}
15171520
}
15181521

15191522
Ok(StageLoadSelectItem {

tests/sqlparser_snowflake.rs

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2363,7 +2363,7 @@ fn test_copy_into_with_transformations() {
23632363
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
23642364
alias: Some(Ident::new("t1")),
23652365
file_col_num: 1,
2366-
element: Some(Ident::new("st")),
2366+
element: Some(vec![Ident::new("st")]),
23672367
item_as: Some(Ident::new("st"))
23682368
})
23692369
);
@@ -2372,7 +2372,7 @@ fn test_copy_into_with_transformations() {
23722372
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
23732373
alias: None,
23742374
file_col_num: 1,
2375-
element: Some(Ident::new("index")),
2375+
element: Some(vec![Ident::new("index")]),
23762376
item_as: None
23772377
})
23782378
);
@@ -2634,6 +2634,28 @@ fn test_snowflake_copy_into_stage_name_ends_with_parens() {
26342634
}
26352635
}
26362636

2637+
#[test]
2638+
fn test_copy_into_with_nested_colon_path() {
2639+
let sql = "COPY INTO tbl (col) FROM (SELECT $1:a:b AS col FROM @stage)";
2640+
match snowflake().verified_stmt(sql) {
2641+
Statement::CopyIntoSnowflake {
2642+
from_transformations,
2643+
..
2644+
} => {
2645+
assert_eq!(
2646+
from_transformations.as_ref().unwrap()[0],
2647+
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
2648+
alias: None,
2649+
file_col_num: 1,
2650+
element: Some(vec![Ident::new("a"), Ident::new("b")]),
2651+
item_as: Some(Ident::new("col"))
2652+
})
2653+
);
2654+
}
2655+
_ => unreachable!(),
2656+
}
2657+
}
2658+
26372659
#[test]
26382660
fn test_snowflake_trim() {
26392661
let real_sql = r#"SELECT customer_id, TRIM(sub_items.value:item_price_id, '"', "a") AS item_price_id FROM models_staging.subscriptions"#;

0 commit comments

Comments
 (0)