diff --git a/src/ast/helpers/stmt_data_loading.rs b/src/ast/helpers/stmt_data_loading.rs index 26bbacc5d1..b37fc87d61 100644 --- a/src/ast/helpers/stmt_data_loading.rs +++ b/src/ast/helpers/stmt_data_loading.rs @@ -135,6 +135,12 @@ pub struct StageLoadSelectItem { pub alias: Option, /// 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, /// Optional element identifier following the column reference. pub element: Option, /// Optional alias for the item (AS clause). @@ -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::().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}")?; } diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs index ad4fd6bb99..76afb78e3a 100644 --- a/src/dialect/snowflake.rs +++ b/src/dialect/snowflake.rs @@ -1773,15 +1773,14 @@ fn parse_select_item_for_data_load( ) -> Result { let mut alias: Option = None; let mut file_col_num: i32 = 0; + let mut file_col_num_raw: Option = None; let mut element: Option = None; let mut item_as: Option = None; let next_token = parser.next_token(); match next_token.token { Token::Placeholder(w) => { - file_col_num = w.to_string().split_off(1).parse::().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) => { @@ -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::().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), @@ -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), ParserError> { + let digits = placeholder.strip_prefix('$').unwrap_or(placeholder); + let number = digits.parse::().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 { let (mut url, mut storage_integration, mut endpoint) = (None, None, None); let mut encryption: KeyValueOptions = KeyValueOptions { diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index 5d7fcb3aaf..859ed7bd1e 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -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")) }) @@ -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 }) @@ -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 }) @@ -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!(