Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion src/ast/helpers/stmt_data_loading.rs
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,12 @@ pub struct StageLoadSelectItem {
pub alias: Option<Ident>,
/// Column number within the staged file (1-based).
pub file_col_num: i32,
/// Original spelling when the staged placeholder differs from its numeric form.
#[cfg_attr(
feature = "serde",
serde(default, skip_serializing_if = "Option::is_none")
)]
pub file_col_num_raw: Option<String>,
/// Optional element identifier following the column reference.
pub element: Option<Ident>,
/// Optional alias for the item (AS clause).
Expand Down Expand Up @@ -172,7 +178,15 @@ impl fmt::Display for StageLoadSelectItem {
if let Some(alias) = &self.alias {
write!(f, "{alias}.")?;
}
write!(f, "${}", self.file_col_num)?;
if let Some(raw) = self.file_col_num_raw.as_deref().filter(|raw| {
raw.strip_prefix('$')
.and_then(|digits| digits.parse::<i32>().ok())
== Some(self.file_col_num)
}) {
write!(f, "{raw}")?;
} else {
write!(f, "${}", self.file_col_num)?;
}
if let Some(element) = &self.element {
write!(f, ":{element}")?;
}
Expand Down
19 changes: 13 additions & 6 deletions src/dialect/snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1773,15 +1773,14 @@ 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 file_col_num_raw: Option<String> = None;
let mut element: Option<Ident> = None;
let mut item_as: Option<Ident> = None;

let next_token = parser.next_token();
match next_token.token {
Token::Placeholder(w) => {
file_col_num = w.to_string().split_off(1).parse::<i32>().map_err(|e| {
ParserError::ParserError(format!("Could not parse '{w}' as i32: {e}"))
})?;
(file_col_num, file_col_num_raw) = parse_stage_load_column_number(&w)?;
Ok(())
}
Token::Word(w) => {
Expand All @@ -1801,9 +1800,7 @@ fn parse_select_item_for_data_load(
let col_num_token = parser.next_token();
match col_num_token.token {
Token::Placeholder(w) => {
file_col_num = w.to_string().split_off(1).parse::<i32>().map_err(|e| {
ParserError::ParserError(format!("Could not parse '{w}' as i32: {e}"))
})?;
(file_col_num, file_col_num_raw) = parse_stage_load_column_number(&w)?;
Ok(())
}
_ => parser.expected("file_col_num", col_num_token),
Expand Down Expand Up @@ -1862,11 +1859,21 @@ fn parse_select_item_for_data_load(
Ok(StageLoadSelectItem {
alias,
file_col_num,
file_col_num_raw,
element,
item_as,
})
}

fn parse_stage_load_column_number(placeholder: &str) -> Result<(i32, Option<String>), ParserError> {
let digits = placeholder.strip_prefix('$').unwrap_or(placeholder);
let number = digits.parse::<i32>().map_err(|error| {
ParserError::ParserError(format!("Could not parse '{placeholder}' as i32: {error}"))
})?;
let raw = (placeholder != format!("${number}")).then(|| String::from(placeholder));
Ok((number, raw))
}

fn parse_stage_params(parser: &mut Parser) -> Result<StageParamsObject, ParserError> {
let (mut url, mut storage_integration, mut endpoint) = (None, None, None);
let mut encryption: KeyValueOptions = KeyValueOptions {
Expand Down
32 changes: 32 additions & 0 deletions tests/sqlparser_snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2703,6 +2703,7 @@ fn test_copy_into_with_transformations() {
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
alias: Some(Ident::new("t1")),
file_col_num: 1,
file_col_num_raw: None,
element: Some(Ident::new("st")),
item_as: Some(Ident::new("st"))
})
Expand All @@ -2712,6 +2713,7 @@ fn test_copy_into_with_transformations() {
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
alias: None,
file_col_num: 1,
file_col_num_raw: None,
element: Some(Ident::new("index")),
item_as: None
})
Expand All @@ -2721,6 +2723,7 @@ fn test_copy_into_with_transformations() {
StageLoadSelectItemKind::StageLoadSelectItem(StageLoadSelectItem {
alias: Some(Ident::new("t2")),
file_col_num: 1,
file_col_num_raw: None,
element: None,
item_as: None
})
Expand Down Expand Up @@ -2782,6 +2785,35 @@ fn test_copy_into_staged_identifiers_preserve_quotes_and_spans() {
assert_eq!(items[0].to_string(), r#""src".$1:"field" AS "Mixed""#);
}

#[test]
fn test_copy_into_staged_placeholders_preserve_zero_padding() {
let sql = r#"COPY INTO t FROM (SELECT $1 AS "$00", $00, src.$01 AS "One" FROM @stage AS src)"#;
let statement = snowflake().verified_stmt(sql);
let Statement::CopyIntoSnowflake {
from_transformations: Some(items),
..
} = statement
else {
panic!("expected COPY INTO with transformations");
};
for (item, number, raw) in [(&items[1], 0, "$00"), (&items[2], 1, "$01")] {
let StageLoadSelectItemKind::StageLoadSelectItem(item) = item else {
panic!("expected a staged field");
};
assert_eq!(item.file_col_num, number);
assert_eq!(item.file_col_num_raw.as_deref(), Some(raw));
}
assert_eq!(items[0].to_string(), "$1 AS \"$00\"");
assert_eq!(items[1].to_string(), "$00");
assert_eq!(items[2].to_string(), "src.$01 AS \"One\"");

let StageLoadSelectItemKind::StageLoadSelectItem(mut changed) = items[2].clone() else {
panic!("expected a staged field");
};
changed.file_col_num = 2;
assert_eq!(changed.to_string(), "src.$2 AS \"One\"");
}

#[test]
fn test_copy_into_with_nested_transformation_path() {
let sql = concat!(
Expand Down
Loading