Snowflake COPY INTO target columns, select items and optional alias (#1805)

yoavcloud · web-flow · commit d090ad4ccfd6 · 2025-04-11T11:49:43.000+02:00
diff --git a/src/ast/helpers/stmt_data_loading.rs b/src/ast/helpers/stmt_data_loading.rs
@@ -29,7 +29,7 @@ use core::fmt;
 use serde::{Deserialize, Serialize};
 
 use crate::ast::helpers::key_value_options::KeyValueOptions;
-use crate::ast::{Ident, ObjectName};
+use crate::ast::{Ident, ObjectName, SelectItem};
 #[cfg(feature = "visitor")]
 use sqlparser_derive::{Visit, VisitMut};
 
@@ -44,6 +44,25 @@ pub struct StageParamsObject {
     pub credentials: KeyValueOptions,
 }
 
+/// This enum enables support for both standard SQL select item expressions
+/// and Snowflake-specific ones for data loading.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum StageLoadSelectItemKind {
+    SelectItem(SelectItem),
+    StageLoadSelectItem(StageLoadSelectItem),
+}
+
+impl fmt::Display for StageLoadSelectItemKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match &self {
+            StageLoadSelectItemKind::SelectItem(item) => write!(f, "{item}"),
+            StageLoadSelectItemKind::StageLoadSelectItem(item) => write!(f, "{item}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -23,7 +23,10 @@ use alloc::{
     string::{String, ToString},
     vec::Vec,
 };
-use helpers::{attached_token::AttachedToken, stmt_data_loading::FileStagingCommand};
+use helpers::{
+    attached_token::AttachedToken,
+    stmt_data_loading::{FileStagingCommand, StageLoadSelectItemKind},
+};
 
 use core::ops::Deref;
 use core::{
@@ -92,7 +95,7 @@ pub use self::value::{
 };
 
 use crate::ast::helpers::key_value_options::KeyValueOptions;
-use crate::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageParamsObject};
+use crate::ast::helpers::stmt_data_loading::StageParamsObject;
 
 #[cfg(feature = "visitor")]
 pub use visitor::*;
@@ -2988,10 +2991,11 @@ pub enum Statement {
     CopyIntoSnowflake {
         kind: CopyIntoSnowflakeKind,
         into: ObjectName,
+        into_columns: Option<Vec<Ident>>,
         from_obj: Option<ObjectName>,
         from_obj_alias: Option<Ident>,
         stage_params: StageParamsObject,
-        from_transformations: Option<Vec<StageLoadSelectItem>>,
+        from_transformations: Option<Vec<StageLoadSelectItemKind>>,
         from_query: Option<Box<Query>>,
         files: Option<Vec<String>>,
         pattern: Option<String>,
@@ -5583,6 +5587,7 @@ impl fmt::Display for Statement {
             Statement::CopyIntoSnowflake {
                 kind,
                 into,
+                into_columns,
                 from_obj,
                 from_obj_alias,
                 stage_params,
@@ -5596,6 +5601,9 @@ impl fmt::Display for Statement {
                 partition,
             } => {
                 write!(f, "COPY INTO {}", into)?;
+                if let Some(into_columns) = into_columns {
+                    write!(f, " ({})", display_comma_separated(into_columns))?;
+                }
                 if let Some(from_transformations) = from_transformations {
                     // Data load with transformation
                     if let Some(from_stage) = from_obj {
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -350,6 +350,7 @@ impl Spanned for Statement {
             } => source.span(),
             Statement::CopyIntoSnowflake {
                 into: _,
+                into_columns: _,
                 from_obj: _,
                 from_obj_alias: _,
                 stage_params: _,
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
@@ -20,7 +20,7 @@ use crate::alloc::string::ToString;
 use crate::ast::helpers::key_value_options::{KeyValueOption, KeyValueOptionType, KeyValueOptions};
 use crate::ast::helpers::stmt_create_table::CreateTableBuilder;
 use crate::ast::helpers::stmt_data_loading::{
-    FileStagingCommand, StageLoadSelectItem, StageParamsObject,
+    FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
 };
 use crate::ast::{
     ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, Ident,
@@ -30,7 +30,7 @@ use crate::ast::{
 };
 use crate::dialect::{Dialect, Precedence};
 use crate::keywords::Keyword;
-use crate::parser::{Parser, ParserError};
+use crate::parser::{IsOptional, Parser, ParserError};
 use crate::tokenizer::{Token, Word};
 #[cfg(not(feature = "std"))]
 use alloc::boxed::Box;
@@ -722,7 +722,7 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
     };
 
     let mut files: Vec<String> = vec![];
-    let mut from_transformations: Option<Vec<StageLoadSelectItem>> = None;
+    let mut from_transformations: Option<Vec<StageLoadSelectItemKind>> = None;
     let mut from_stage_alias = None;
     let mut from_stage = None;
     let mut stage_params = StageParamsObject {
@@ -744,6 +744,11 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
         stage_params = parse_stage_params(parser)?;
     }
 
+    let into_columns = match &parser.peek_token().token {
+        Token::LParen => Some(parser.parse_parenthesized_column_list(IsOptional::Optional, true)?),
+        _ => None,
+    };
+
     parser.expect_keyword_is(Keyword::FROM)?;
     match parser.next_token().token {
         Token::LParen if kind == CopyIntoSnowflakeKind::Table => {
@@ -755,15 +760,10 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
             from_stage = Some(parse_snowflake_stage_name(parser)?);
             stage_params = parse_stage_params(parser)?;
 
-            // as
-            from_stage_alias = if parser.parse_keyword(Keyword::AS) {
-                Some(match parser.next_token().token {
-                    Token::Word(w) => Ok(Ident::new(w.value)),
-                    _ => parser.expected("stage alias", parser.peek_token()),
-                }?)
-            } else {
-                None
-            };
+            // Parse an optional alias
+            from_stage_alias = parser
+                .maybe_parse_table_alias()?
+                .map(|table_alias| table_alias.name);
             parser.expect_token(&Token::RParen)?;
         }
         Token::LParen if kind == CopyIntoSnowflakeKind::Location => {
@@ -846,6 +846,7 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
     Ok(Statement::CopyIntoSnowflake {
         kind,
         into,
+        into_columns,
         from_obj: from_stage,
         from_obj_alias: from_stage_alias,
         stage_params,
@@ -866,86 +867,93 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
 
 fn parse_select_items_for_data_load(
     parser: &mut Parser,
-) -> Result<Option<Vec<StageLoadSelectItem>>, ParserError> {
-    // [<alias>.]$<file_col_num>[.<element>] [ , [<alias>.]$<file_col_num>[.<element>] ... ]
-    let mut select_items: Vec<StageLoadSelectItem> = vec![];
+) -> Result<Option<Vec<StageLoadSelectItemKind>>, ParserError> {
+    let mut select_items: Vec<StageLoadSelectItemKind> = vec![];
     loop {
-        let mut alias: Option<Ident> = None;
-        let mut file_col_num: i32 = 0;
-        let mut element: Option<Ident> = None;
-        let mut item_as: Option<Ident> = None;
+        match parser.maybe_parse(parse_select_item_for_data_load)? {
+            // [<alias>.]$<file_col_num>[.<element>] [ , [<alias>.]$<file_col_num>[.<element>] ... ]
+            Some(item) => select_items.push(StageLoadSelectItemKind::StageLoadSelectItem(item)),
+            // Fallback, try to parse a standard SQL select item
+            None => select_items.push(StageLoadSelectItemKind::SelectItem(
+                parser.parse_select_item()?,
+            )),
+        }
+        if matches!(parser.peek_token_ref().token, Token::Comma) {
+            parser.advance_token();
+        } else {
+            break;
+        }
+    }
+    Ok(Some(select_items))
+}
 
-        let next_token = parser.next_token();
-        match next_token.token {
+fn parse_select_item_for_data_load(
+    parser: &mut Parser,
+) -> Result<StageLoadSelectItem, ParserError> {
+    let mut alias: Option<Ident> = None;
+    let mut file_col_num: i32 = 0;
+    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}"))
+            })?;
+            Ok(())
+        }
+        Token::Word(w) => {
+            alias = Some(Ident::new(w.value));
+            Ok(())
+        }
+        _ => parser.expected("alias or file_col_num", next_token),
+    }?;
+
+    if alias.is_some() {
+        parser.expect_token(&Token::Period)?;
+        // now we get col_num token
+        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}"))
                 })?;
                 Ok(())
             }
-            Token::Word(w) => {
-                alias = Some(Ident::new(w.value));
-                Ok(())
-            }
-            _ => parser.expected("alias or file_col_num", next_token),
+            _ => parser.expected("file_col_num", col_num_token),
         }?;
+    }
 
-        if alias.is_some() {
-            parser.expect_token(&Token::Period)?;
-            // now we get col_num token
-            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}"))
-                    })?;
-                    Ok(())
-                }
-                _ => parser.expected("file_col_num", col_num_token),
-            }?;
-        }
-
-        // 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
+    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()),
+            }?));
         }
-
-        // as
-        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()),
-            }?);
+        _ => {
+            // element not present move back
+            parser.prev_token();
         }
+    }
 
-        select_items.push(StageLoadSelectItem {
-            alias,
-            file_col_num,
-            element,
-            item_as,
-        });
-
-        match parser.next_token().token {
-            Token::Comma => {
-                // continue
-            }
-            _ => {
-                parser.prev_token(); // need to move back
-                break;
-            }
-        }
+    // as
+    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()),
+        }?);
     }
-    Ok(Some(select_items))
+
+    Ok(StageLoadSelectItem {
+        alias,
+        file_col_num,
+        element,
+        item_as,
+    })
 }
 
 fn parse_stage_params(parser: &mut Parser) -> Result<StageParamsObject, ParserError> {
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
