Preserve zero-padded COPY staged placeholders

Author: osipovartem

src/ast/helpers/stmt_data_loading.rs

@@ -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).
@@ -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}")?;
         }

src/dialect/snowflake.rs

@@ -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) => {
@@ -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),
@@ -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(|| placeholder.to_owned());
+    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 {

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!(