Skip to content

Commit d090ad4

Browse files
yoavcloudyoavcloud
authored
Snowflake COPY INTO target columns, select items and optional alias (apache#1805)
1 parent 67c3be0 commit d090ad4

File tree

5 files changed

+146
-90
lines changed

5 files changed

+146
-90
lines changed

‎src/ast/helpers/stmt_data_loading.rs‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use core::fmt;
2929
use serde::{Deserialize, Serialize};
3030

3131
use crate::ast::helpers::key_value_options::KeyValueOptions;
32-
use crate::ast::{Ident, ObjectName};
32+
use crate::ast::{Ident, ObjectName, SelectItem};
3333
#[cfg(feature = "visitor")]
3434
use sqlparser_derive::{Visit, VisitMut};
3535

@@ -44,6 +44,25 @@ pub struct StageParamsObject {
4444
pub credentials: KeyValueOptions,
4545
}
4646

47+
/// This enum enables support for both standard SQL select item expressions
48+
/// and Snowflake-specific ones for data loading.
49+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
50+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
51+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
52+
pub enum StageLoadSelectItemKind {
53+
SelectItem(SelectItem),
54+
StageLoadSelectItem(StageLoadSelectItem),
55+
}
56+
57+
impl fmt::Display for StageLoadSelectItemKind {
58+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
59+
match &self {
60+
StageLoadSelectItemKind::SelectItem(item) => write!(f, "{item}"),
61+
StageLoadSelectItemKind::StageLoadSelectItem(item) => write!(f, "{item}"),
62+
}
63+
}
64+
}
65+
4766
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
4867
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
4968
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

‎src/ast/mod.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,10 @@ use alloc::{
2323
string::{String, ToString},
2424
vec::Vec,
2525
};
26-
use helpers::{attached_token::AttachedToken, stmt_data_loading::FileStagingCommand};
26+
use helpers::{
27+
attached_token::AttachedToken,
28+
stmt_data_loading::{FileStagingCommand, StageLoadSelectItemKind},
29+
};
2730

2831
use core::ops::Deref;
2932
use core::{
@@ -92,7 +95,7 @@ pub use self::value::{
9295
};
9396

9497
use crate::ast::helpers::key_value_options::KeyValueOptions;
95-
use crate::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageParamsObject};
98+
use crate::ast::helpers::stmt_data_loading::StageParamsObject;
9699

97100
#[cfg(feature = "visitor")]
98101
pub use visitor::*;
@@ -2988,10 +2991,11 @@ pub enum Statement {
29882991
CopyIntoSnowflake {
29892992
kind: CopyIntoSnowflakeKind,
29902993
into: ObjectName,
2994+
into_columns: Option<Vec<Ident>>,
29912995
from_obj: Option<ObjectName>,
29922996
from_obj_alias: Option<Ident>,
29932997
stage_params: StageParamsObject,
2994-
from_transformations: Option<Vec<StageLoadSelectItem>>,
2998+
from_transformations: Option<Vec<StageLoadSelectItemKind>>,
29952999
from_query: Option<Box<Query>>,
29963000
files: Option<Vec<String>>,
29973001
pattern: Option<String>,
@@ -5583,6 +5587,7 @@ impl fmt::Display for Statement {
55835587
Statement::CopyIntoSnowflake {
55845588
kind,
55855589
into,
5590+
into_columns,
55865591
from_obj,
55875592
from_obj_alias,
55885593
stage_params,
@@ -5596,6 +5601,9 @@ impl fmt::Display for Statement {
55965601
partition,
55975602
} => {
55985603
write!(f, "COPY INTO {}", into)?;
5604+
if let Some(into_columns) = into_columns {
5605+
write!(f, " ({})", display_comma_separated(into_columns))?;
5606+
}
55995607
if let Some(from_transformations) = from_transformations {
56005608
// Data load with transformation
56015609
if let Some(from_stage) = from_obj {

‎src/ast/spans.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,7 @@ impl Spanned for Statement {
350350
} => source.span(),
351351
Statement::CopyIntoSnowflake {
352352
into: _,
353+
into_columns: _,
353354
from_obj: _,
354355
from_obj_alias: _,
355356
stage_params: _,

‎src/dialect/snowflake.rs‎

Lines changed: 86 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use crate::alloc::string::ToString;
2020
use crate::ast::helpers::key_value_options::{KeyValueOption, KeyValueOptionType, KeyValueOptions};
2121
use crate::ast::helpers::stmt_create_table::CreateTableBuilder;
2222
use crate::ast::helpers::stmt_data_loading::{
23-
FileStagingCommand, StageLoadSelectItem, StageParamsObject,
23+
FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
2424
};
2525
use crate::ast::{
2626
ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind, Ident,
@@ -30,7 +30,7 @@ use crate::ast::{
3030
};
3131
use crate::dialect::{Dialect, Precedence};
3232
use crate::keywords::Keyword;
33-
use crate::parser::{Parser, ParserError};
33+
use crate::parser::{IsOptional, Parser, ParserError};
3434
use crate::tokenizer::{Token, Word};
3535
#[cfg(not(feature = "std"))]
3636
use alloc::boxed::Box;
@@ -722,7 +722,7 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
722722
};
723723

724724
let mut files: Vec<String> = vec![];
725-
let mut from_transformations: Option<Vec<StageLoadSelectItem>> = None;
725+
let mut from_transformations: Option<Vec<StageLoadSelectItemKind>> = None;
726726
let mut from_stage_alias = None;
727727
let mut from_stage = None;
728728
let mut stage_params = StageParamsObject {
@@ -744,6 +744,11 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
744744
stage_params = parse_stage_params(parser)?;
745745
}
746746

747+
let into_columns = match &parser.peek_token().token {
748+
Token::LParen => Some(parser.parse_parenthesized_column_list(IsOptional::Optional, true)?),
749+
_ => None,
750+
};
751+
747752
parser.expect_keyword_is(Keyword::FROM)?;
748753
match parser.next_token().token {
749754
Token::LParen if kind == CopyIntoSnowflakeKind::Table => {
@@ -755,15 +760,10 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
755760
from_stage = Some(parse_snowflake_stage_name(parser)?);
756761
stage_params = parse_stage_params(parser)?;
757762

758-
// as
759-
from_stage_alias = if parser.parse_keyword(Keyword::AS) {
760-
Some(match parser.next_token().token {
761-
Token::Word(w) => Ok(Ident::new(w.value)),
762-
_ => parser.expected("stage alias", parser.peek_token()),
763-
}?)
764-
} else {
765-
None
766-
};
763+
// Parse an optional alias
764+
from_stage_alias = parser
765+
.maybe_parse_table_alias()?
766+
.map(|table_alias| table_alias.name);
767767
parser.expect_token(&Token::RParen)?;
768768
}
769769
Token::LParen if kind == CopyIntoSnowflakeKind::Location => {
@@ -846,6 +846,7 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
846846
Ok(Statement::CopyIntoSnowflake {
847847
kind,
848848
into,
849+
into_columns,
849850
from_obj: from_stage,
850851
from_obj_alias: from_stage_alias,
851852
stage_params,
@@ -866,86 +867,93 @@ pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
866867

867868
fn parse_select_items_for_data_load(
868869
parser: &mut Parser,
869-
) -> Result<Option<Vec<StageLoadSelectItem>>, ParserError> {
870-
// [<alias>.]$<file_col_num>[.<element>] [ , [<alias>.]$<file_col_num>[.<element>] ... ]
871-
let mut select_items: Vec<StageLoadSelectItem> = vec![];
870+
) -> Result<Option<Vec<StageLoadSelectItemKind>>, ParserError> {
871+
let mut select_items: Vec<StageLoadSelectItemKind> = vec![];
872872
loop {
873-
let mut alias: Option<Ident> = None;
874-
let mut file_col_num: i32 = 0;
875-
let mut element: Option<Ident> = None;
876-
let mut item_as: Option<Ident> = None;
873+
match parser.maybe_parse(parse_select_item_for_data_load)? {
874+
// [<alias>.]$<file_col_num>[.<element>] [ , [<alias>.]$<file_col_num>[.<element>] ... ]
875+
Some(item) => select_items.push(StageLoadSelectItemKind::StageLoadSelectItem(item)),
876+
// Fallback, try to parse a standard SQL select item
877+
None => select_items.push(StageLoadSelectItemKind::SelectItem(
878+
parser.parse_select_item()?,
879+
)),
880+
}
881+
if matches!(parser.peek_token_ref().token, Token::Comma) {
882+
parser.advance_token();
883+
} else {
884+
break;
885+
}
886+
}
887+
Ok(Some(select_items))
888+
}
877889

878-
let next_token = parser.next_token();
879-
match next_token.token {
890+
fn parse_select_item_for_data_load(
891+
parser: &mut Parser,
892+
) -> Result<StageLoadSelectItem, ParserError> {
893+
let mut alias: Option<Ident> = None;
894+
let mut file_col_num: i32 = 0;
895+
let mut element: Option<Ident> = None;
896+
let mut item_as: Option<Ident> = None;
897+
898+
let next_token = parser.next_token();
899+
match next_token.token {
900+
Token::Placeholder(w) => {
901+
file_col_num = w.to_string().split_off(1).parse::<i32>().map_err(|e| {
902+
ParserError::ParserError(format!("Could not parse '{w}' as i32: {e}"))
903+
})?;
904+
Ok(())
905+
}
906+
Token::Word(w) => {
907+
alias = Some(Ident::new(w.value));
908+
Ok(())
909+
}
910+
_ => parser.expected("alias or file_col_num", next_token),
911+
}?;
912+
913+
if alias.is_some() {
914+
parser.expect_token(&Token::Period)?;
915+
// now we get col_num token
916+
let col_num_token = parser.next_token();
917+
match col_num_token.token {
880918
Token::Placeholder(w) => {
881919
file_col_num = w.to_string().split_off(1).parse::<i32>().map_err(|e| {
882920
ParserError::ParserError(format!("Could not parse '{w}' as i32: {e}"))
883921
})?;
884922
Ok(())
885923
}
886-
Token::Word(w) => {
887-
alias = Some(Ident::new(w.value));
888-
Ok(())
889-
}
890-
_ => parser.expected("alias or file_col_num", next_token),
924+
_ => parser.expected("file_col_num", col_num_token),
891925
}?;
926+
}
892927

893-
if alias.is_some() {
894-
parser.expect_token(&Token::Period)?;
895-
// now we get col_num token
896-
let col_num_token = parser.next_token();
897-
match col_num_token.token {
898-
Token::Placeholder(w) => {
899-
file_col_num = w.to_string().split_off(1).parse::<i32>().map_err(|e| {
900-
ParserError::ParserError(format!("Could not parse '{w}' as i32: {e}"))
901-
})?;
902-
Ok(())
903-
}
904-
_ => parser.expected("file_col_num", col_num_token),
905-
}?;
906-
}
907-
908-
// try extracting optional element
909-
match parser.next_token().token {
910-
Token::Colon => {
911-
// parse element
912-
element = Some(Ident::new(match parser.next_token().token {
913-
Token::Word(w) => Ok(w.value),
914-
_ => parser.expected("file_col_num", parser.peek_token()),
915-
}?));
916-
}
917-
_ => {
918-
// element not present move back
919-
parser.prev_token();
920-
}
928+
// try extracting optional element
929+
match parser.next_token().token {
930+
Token::Colon => {
931+
// parse element
932+
element = Some(Ident::new(match parser.next_token().token {
933+
Token::Word(w) => Ok(w.value),
934+
_ => parser.expected("file_col_num", parser.peek_token()),
935+
}?));
921936
}
922-
923-
// as
924-
if parser.parse_keyword(Keyword::AS) {
925-
item_as = Some(match parser.next_token().token {
926-
Token::Word(w) => Ok(Ident::new(w.value)),
927-
_ => parser.expected("column item alias", parser.peek_token()),
928-
}?);
937+
_ => {
938+
// element not present move back
939+
parser.prev_token();
929940
}
941+
}
930942

931-
select_items.push(StageLoadSelectItem {
932-
alias,
933-
file_col_num,
934-
element,
935-
item_as,
936-
});
937-
938-
match parser.next_token().token {
939-
Token::Comma => {
940-
// continue
941-
}
942-
_ => {
943-
parser.prev_token(); // need to move back
944-
break;
945-
}
946-
}
943+
// as
944+
if parser.parse_keyword(Keyword::AS) {
945+
item_as = Some(match parser.next_token().token {
946+
Token::Word(w) => Ok(Ident::new(w.value)),
947+
_ => parser.expected("column item alias", parser.peek_token()),
948+
}?);
947949
}
948-
Ok(Some(select_items))
950+
951+
Ok(StageLoadSelectItem {
952+
alias,
953+
file_col_num,
954+
element,
955+
item_as,
956+
})
949957
}
950958

951959
fn parse_stage_params(parser: &mut Parser) -> Result<StageParamsObject, ParserError> {

0 commit comments

Comments
 (0)