This is an automated email from the ASF dual-hosted git repository.
iffyio pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-sqlparser-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 751dc5af Parse Snowflake COPY INTO <location> (#1669)
751dc5af is described below
commit 751dc5afce794398bcfc148aa6b2f9ecf3b6b5a9
Author: Yoav Cohen <[email protected]>
AuthorDate: Wed Feb 5 20:23:27 2025 +0100
Parse Snowflake COPY INTO <location> (#1669)
---
src/ast/helpers/stmt_data_loading.rs | 10 +-
src/ast/mod.rs | 102 ++++++++------
src/ast/spans.rs | 7 +-
src/dialect/snowflake.rs | 250 ++++++++++++++++++++---------------
tests/sqlparser_snowflake.rs | 197 ++++++++++++++++++++++-----
5 files changed, 381 insertions(+), 185 deletions(-)
diff --git a/src/ast/helpers/stmt_data_loading.rs
b/src/ast/helpers/stmt_data_loading.rs
index 42e1df06..77de5d9e 100644
--- a/src/ast/helpers/stmt_data_loading.rs
+++ b/src/ast/helpers/stmt_data_loading.rs
@@ -58,6 +58,7 @@ pub enum DataLoadingOptionType {
STRING,
BOOLEAN,
ENUM,
+ NUMBER,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -128,12 +129,9 @@ impl fmt::Display for DataLoadingOption {
DataLoadingOptionType::STRING => {
write!(f, "{}='{}'", self.option_name, self.value)?;
}
- DataLoadingOptionType::ENUM => {
- // single quote is omitted
- write!(f, "{}={}", self.option_name, self.value)?;
- }
- DataLoadingOptionType::BOOLEAN => {
- // single quote is omitted
+ DataLoadingOptionType::ENUM
+ | DataLoadingOptionType::BOOLEAN
+ | DataLoadingOptionType::NUMBER => {
write!(f, "{}={}", self.option_name, self.value)?;
}
}
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 35c6dcc1..dc944c9e 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -2498,24 +2498,30 @@ pub enum Statement {
values: Vec<Option<String>>,
},
/// ```sql
- /// COPY INTO
+ /// COPY INTO <table> | <location>
/// ```
- /// See <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+ /// See:
+ /// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+ /// <https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
+ ///
/// Copy Into syntax available for Snowflake is different than the one
implemented in
/// Postgres. Although they share common prefix, it is reasonable to
implement them
/// in different enums. This can be refactored later once custom dialects
/// are allowed to have custom Statements.
CopyIntoSnowflake {
+ kind: CopyIntoSnowflakeKind,
into: ObjectName,
- from_stage: ObjectName,
- from_stage_alias: Option<Ident>,
+ from_obj: Option<ObjectName>,
+ from_obj_alias: Option<Ident>,
stage_params: StageParamsObject,
from_transformations: Option<Vec<StageLoadSelectItem>>,
+ from_query: Option<Box<Query>>,
files: Option<Vec<String>>,
pattern: Option<String>,
file_format: DataLoadingOptions,
copy_options: DataLoadingOptions,
validation_mode: Option<String>,
+ partition: Option<Box<Expr>>,
},
/// ```sql
/// CLOSE
@@ -5048,60 +5054,69 @@ impl fmt::Display for Statement {
Ok(())
}
Statement::CopyIntoSnowflake {
+ kind,
into,
- from_stage,
- from_stage_alias,
+ from_obj,
+ from_obj_alias,
stage_params,
from_transformations,
+ from_query,
files,
pattern,
file_format,
copy_options,
validation_mode,
+ partition,
} => {
write!(f, "COPY INTO {}", into)?;
- if from_transformations.is_none() {
- // Standard data load
- write!(f, " FROM {}{}", from_stage, stage_params)?;
- if from_stage_alias.as_ref().is_some() {
- write!(f, " AS {}",
from_stage_alias.as_ref().unwrap())?;
- }
- } else {
+ if let Some(from_transformations) = from_transformations {
// Data load with transformation
- write!(
- f,
- " FROM (SELECT {} FROM {}{}",
-
display_separated(from_transformations.as_ref().unwrap(), ", "),
- from_stage,
- stage_params,
- )?;
- if from_stage_alias.as_ref().is_some() {
- write!(f, " AS {}",
from_stage_alias.as_ref().unwrap())?;
+ if let Some(from_stage) = from_obj {
+ write!(
+ f,
+ " FROM (SELECT {} FROM {}{}",
+ display_separated(from_transformations, ", "),
+ from_stage,
+ stage_params
+ )?;
+ }
+ if let Some(from_obj_alias) = from_obj_alias {
+ write!(f, " AS {}", from_obj_alias)?;
}
write!(f, ")")?;
+ } else if let Some(from_obj) = from_obj {
+ // Standard data load
+ write!(f, " FROM {}{}", from_obj, stage_params)?;
+ if let Some(from_obj_alias) = from_obj_alias {
+ write!(f, " AS {from_obj_alias}")?;
+ }
+ } else if let Some(from_query) = from_query {
+ // Data unload from query
+ write!(f, " FROM ({from_query})")?;
}
- if files.is_some() {
- write!(
- f,
- " FILES = ('{}')",
- display_separated(files.as_ref().unwrap(), "', '")
- )?;
+
+ if let Some(files) = files {
+ write!(f, " FILES = ('{}')", display_separated(files, "',
'"))?;
+ }
+ if let Some(pattern) = pattern {
+ write!(f, " PATTERN = '{}'", pattern)?;
}
- if pattern.is_some() {
- write!(f, " PATTERN = '{}'", pattern.as_ref().unwrap())?;
+ if let Some(partition) = partition {
+ write!(f, " PARTITION BY {partition}")?;
}
if !file_format.options.is_empty() {
write!(f, " FILE_FORMAT=({})", file_format)?;
}
if !copy_options.options.is_empty() {
- write!(f, " COPY_OPTIONS=({})", copy_options)?;
+ match kind {
+ CopyIntoSnowflakeKind::Table => {
+ write!(f, " COPY_OPTIONS=({})", copy_options)?
+ }
+ CopyIntoSnowflakeKind::Location => write!(f, "
{copy_options}")?,
+ }
}
- if validation_mode.is_some() {
- write!(
- f,
- " VALIDATION_MODE = {}",
- validation_mode.as_ref().unwrap()
- )?;
+ if let Some(validation_mode) = validation_mode {
+ write!(f, " VALIDATION_MODE = {}", validation_mode)?;
}
Ok(())
}
@@ -8543,6 +8558,19 @@ impl Display for StorageSerializationPolicy {
}
}
+/// Variants of the Snowflake `COPY INTO` statement
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum CopyIntoSnowflakeKind {
+ /// Loads data from files to a table
+ /// See: <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table>
+ Table,
+ /// Unloads data from a table or query to external files
+ /// See:
<https://docs.snowflake.com/en/sql-reference/sql/copy-into-location>
+ Location,
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
index b2690085..f0c38942 100644
--- a/src/ast/spans.rs
+++ b/src/ast/spans.rs
@@ -333,8 +333,8 @@ impl Spanned for Statement {
} => source.span(),
Statement::CopyIntoSnowflake {
into: _,
- from_stage: _,
- from_stage_alias: _,
+ from_obj: _,
+ from_obj_alias: _,
stage_params: _,
from_transformations: _,
files: _,
@@ -342,6 +342,9 @@ impl Spanned for Statement {
file_format: _,
copy_options: _,
validation_mode: _,
+ kind: _,
+ from_query: _,
+ partition: _,
} => Span::empty(),
Statement::Close { cursor } => match cursor {
CloseCursor::All => Span::empty(),
diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs
index d775ffc3..fc192671 100644
--- a/src/dialect/snowflake.rs
+++ b/src/dialect/snowflake.rs
@@ -23,23 +23,26 @@ use crate::ast::helpers::stmt_data_loading::{
StageLoadSelectItem, StageParamsObject,
};
use crate::ast::{
- ColumnOption, ColumnPolicy, ColumnPolicyProperty, Ident,
IdentityParameters, IdentityProperty,
- IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder,
ObjectName,
- RowAccessPolicy, Statement, TagsColumnOption, WrappedCollection,
+ ColumnOption, ColumnPolicy, ColumnPolicyProperty, CopyIntoSnowflakeKind,
Ident,
+ IdentityParameters, IdentityProperty, IdentityPropertyFormatKind,
IdentityPropertyKind,
+ IdentityPropertyOrder, ObjectName, RowAccessPolicy, Statement,
TagsColumnOption,
+ WrappedCollection,
};
use crate::dialect::{Dialect, Precedence};
use crate::keywords::Keyword;
use crate::parser::{Parser, ParserError};
-use crate::tokenizer::Token;
+use crate::tokenizer::{Token, Word};
+#[cfg(not(feature = "std"))]
+use alloc::boxed::Box;
#[cfg(not(feature = "std"))]
use alloc::string::String;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(not(feature = "std"))]
use alloc::{format, vec};
-use sqlparser::ast::StorageSerializationPolicy;
use super::keywords::RESERVED_FOR_IDENTIFIER;
+use sqlparser::ast::StorageSerializationPolicy;
/// A [`Dialect`] for [Snowflake](https://www.snowflake.com/)
#[derive(Debug, Default)]
@@ -665,24 +668,49 @@ pub fn parse_snowflake_stage_name(parser: &mut Parser) ->
Result<ObjectName, Par
}
}
+/// Parses a `COPY INTO` statement. Snowflake has two variants, `COPY INTO
<table>`
+/// and `COPY INTO <location>` which have different syntax.
pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
- let into: ObjectName = parse_snowflake_stage_name(parser)?;
+ let kind = match parser.peek_token().token {
+ // Indicates an internal stage
+ Token::AtSign => CopyIntoSnowflakeKind::Location,
+ // Indicates an external stage, i.e. s3://, gcs:// or azure://
+ Token::SingleQuotedString(s) if s.contains("://") =>
CopyIntoSnowflakeKind::Location,
+ _ => CopyIntoSnowflakeKind::Table,
+ };
+
let mut files: Vec<String> = vec![];
let mut from_transformations: Option<Vec<StageLoadSelectItem>> = None;
- let from_stage_alias;
- let from_stage: ObjectName;
- let stage_params: StageParamsObject;
+ let mut from_stage_alias = None;
+ let mut from_stage = None;
+ let mut stage_params = StageParamsObject {
+ url: None,
+ encryption: DataLoadingOptions { options: vec![] },
+ endpoint: None,
+ storage_integration: None,
+ credentials: DataLoadingOptions { options: vec![] },
+ };
+ let mut from_query = None;
+ let mut partition = None;
+ let mut file_format = Vec::new();
+ let mut pattern = None;
+ let mut validation_mode = None;
+ let mut copy_options = Vec::new();
+
+ let into: ObjectName = parse_snowflake_stage_name(parser)?;
+ if kind == CopyIntoSnowflakeKind::Location {
+ stage_params = parse_stage_params(parser)?;
+ }
parser.expect_keyword_is(Keyword::FROM)?;
- // check if data load transformations are present
match parser.next_token().token {
- Token::LParen => {
- // data load with transformations
+ Token::LParen if kind == CopyIntoSnowflakeKind::Table => {
+ // Data load with transformations
parser.expect_keyword_is(Keyword::SELECT)?;
from_transformations = parse_select_items_for_data_load(parser)?;
parser.expect_keyword_is(Keyword::FROM)?;
- from_stage = parse_snowflake_stage_name(parser)?;
+ from_stage = Some(parse_snowflake_stage_name(parser)?);
stage_params = parse_stage_params(parser)?;
// as
@@ -696,9 +724,14 @@ pub fn parse_copy_into(parser: &mut Parser) ->
Result<Statement, ParserError> {
};
parser.expect_token(&Token::RParen)?;
}
+ Token::LParen if kind == CopyIntoSnowflakeKind::Location => {
+ // Data unload with a query
+ from_query = Some(parser.parse_query()?);
+ parser.expect_token(&Token::RParen)?;
+ }
_ => {
parser.prev_token();
- from_stage = parse_snowflake_stage_name(parser)?;
+ from_stage = Some(parse_snowflake_stage_name(parser)?);
stage_params = parse_stage_params(parser)?;
// as
@@ -711,67 +744,71 @@ pub fn parse_copy_into(parser: &mut Parser) ->
Result<Statement, ParserError> {
None
};
}
- };
+ }
- // [ files ]
- if parser.parse_keyword(Keyword::FILES) {
- parser.expect_token(&Token::Eq)?;
- parser.expect_token(&Token::LParen)?;
- let mut continue_loop = true;
- while continue_loop {
- continue_loop = false;
+ loop {
+ // FILE_FORMAT
+ if parser.parse_keyword(Keyword::FILE_FORMAT) {
+ parser.expect_token(&Token::Eq)?;
+ file_format = parse_parentheses_options(parser)?;
+ // PARTITION BY
+ } else if parser.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) {
+ partition = Some(Box::new(parser.parse_expr()?))
+ // FILES
+ } else if parser.parse_keyword(Keyword::FILES) {
+ parser.expect_token(&Token::Eq)?;
+ parser.expect_token(&Token::LParen)?;
+ let mut continue_loop = true;
+ while continue_loop {
+ continue_loop = false;
+ let next_token = parser.next_token();
+ match next_token.token {
+ Token::SingleQuotedString(s) => files.push(s),
+ _ => parser.expected("file token", next_token)?,
+ };
+ if parser.next_token().token.eq(&Token::Comma) {
+ continue_loop = true;
+ } else {
+ parser.prev_token(); // not a comma, need to go back
+ }
+ }
+ parser.expect_token(&Token::RParen)?;
+ // PATTERN
+ } else if parser.parse_keyword(Keyword::PATTERN) {
+ parser.expect_token(&Token::Eq)?;
let next_token = parser.next_token();
- match next_token.token {
- Token::SingleQuotedString(s) => files.push(s),
- _ => parser.expected("file token", next_token)?,
- };
- if parser.next_token().token.eq(&Token::Comma) {
- continue_loop = true;
- } else {
- parser.prev_token(); // not a comma, need to go back
+ pattern = Some(match next_token.token {
+ Token::SingleQuotedString(s) => s,
+ _ => parser.expected("pattern", next_token)?,
+ });
+ // VALIDATION MODE
+ } else if parser.parse_keyword(Keyword::VALIDATION_MODE) {
+ parser.expect_token(&Token::Eq)?;
+ validation_mode = Some(parser.next_token().token.to_string());
+ // COPY OPTIONS
+ } else if parser.parse_keyword(Keyword::COPY_OPTIONS) {
+ parser.expect_token(&Token::Eq)?;
+ copy_options = parse_parentheses_options(parser)?;
+ } else {
+ match parser.next_token().token {
+ Token::SemiColon | Token::EOF => break,
+ Token::Comma => continue,
+ // In `COPY INTO <location>` the copy options do not have a
shared key
+ // like in `COPY INTO <table>`
+ Token::Word(key) =>
copy_options.push(parse_copy_option(parser, key)?),
+ _ => return parser.expected("another copy option, ; or EOF'",
parser.peek_token()),
}
}
- parser.expect_token(&Token::RParen)?;
- }
-
- // [ pattern ]
- let mut pattern = None;
- if parser.parse_keyword(Keyword::PATTERN) {
- parser.expect_token(&Token::Eq)?;
- let next_token = parser.next_token();
- pattern = Some(match next_token.token {
- Token::SingleQuotedString(s) => s,
- _ => parser.expected("pattern", next_token)?,
- });
- }
-
- // [ file_format]
- let mut file_format = Vec::new();
- if parser.parse_keyword(Keyword::FILE_FORMAT) {
- parser.expect_token(&Token::Eq)?;
- file_format = parse_parentheses_options(parser)?;
- }
-
- // [ copy_options ]
- let mut copy_options = Vec::new();
- if parser.parse_keyword(Keyword::COPY_OPTIONS) {
- parser.expect_token(&Token::Eq)?;
- copy_options = parse_parentheses_options(parser)?;
- }
-
- // [ VALIDATION_MODE ]
- let mut validation_mode = None;
- if parser.parse_keyword(Keyword::VALIDATION_MODE) {
- parser.expect_token(&Token::Eq)?;
- validation_mode = Some(parser.next_token().token.to_string());
}
Ok(Statement::CopyIntoSnowflake {
+ kind,
into,
- from_stage,
- from_stage_alias,
+ from_obj: from_stage,
+ from_obj_alias: from_stage_alias,
stage_params,
from_transformations,
+ from_query,
files: if files.is_empty() { None } else { Some(files) },
pattern,
file_format: DataLoadingOptions {
@@ -781,6 +818,7 @@ pub fn parse_copy_into(parser: &mut Parser) ->
Result<Statement, ParserError> {
options: copy_options,
},
validation_mode,
+ partition,
})
}
@@ -930,55 +968,55 @@ fn parse_stage_params(parser: &mut Parser) ->
Result<StageParamsObject, ParserEr
///
fn parse_parentheses_options(parser: &mut Parser) ->
Result<Vec<DataLoadingOption>, ParserError> {
let mut options: Vec<DataLoadingOption> = Vec::new();
-
parser.expect_token(&Token::LParen)?;
loop {
match parser.next_token().token {
Token::RParen => break,
- Token::Word(key) => {
- parser.expect_token(&Token::Eq)?;
- if parser.parse_keyword(Keyword::TRUE) {
- options.push(DataLoadingOption {
- option_name: key.value,
- option_type: DataLoadingOptionType::BOOLEAN,
- value: "TRUE".to_string(),
- });
- Ok(())
- } else if parser.parse_keyword(Keyword::FALSE) {
- options.push(DataLoadingOption {
- option_name: key.value,
- option_type: DataLoadingOptionType::BOOLEAN,
- value: "FALSE".to_string(),
- });
- Ok(())
- } else {
- match parser.next_token().token {
- Token::SingleQuotedString(value) => {
- options.push(DataLoadingOption {
- option_name: key.value,
- option_type: DataLoadingOptionType::STRING,
- value,
- });
- Ok(())
- }
- Token::Word(word) => {
- options.push(DataLoadingOption {
- option_name: key.value,
- option_type: DataLoadingOptionType::ENUM,
- value: word.value,
- });
- Ok(())
- }
- _ => parser.expected("expected option value",
parser.peek_token()),
- }
- }
- }
- _ => parser.expected("another option or ')'", parser.peek_token()),
- }?;
+ Token::Comma => continue,
+ Token::Word(key) => options.push(parse_copy_option(parser, key)?),
+ _ => return parser.expected("another option or ')'",
parser.peek_token()),
+ };
}
Ok(options)
}
+/// Parses a `KEY = VALUE` construct based on the specified key
+fn parse_copy_option(parser: &mut Parser, key: Word) ->
Result<DataLoadingOption, ParserError> {
+ parser.expect_token(&Token::Eq)?;
+ if parser.parse_keyword(Keyword::TRUE) {
+ Ok(DataLoadingOption {
+ option_name: key.value,
+ option_type: DataLoadingOptionType::BOOLEAN,
+ value: "TRUE".to_string(),
+ })
+ } else if parser.parse_keyword(Keyword::FALSE) {
+ Ok(DataLoadingOption {
+ option_name: key.value,
+ option_type: DataLoadingOptionType::BOOLEAN,
+ value: "FALSE".to_string(),
+ })
+ } else {
+ match parser.next_token().token {
+ Token::SingleQuotedString(value) => Ok(DataLoadingOption {
+ option_name: key.value,
+ option_type: DataLoadingOptionType::STRING,
+ value,
+ }),
+ Token::Word(word) => Ok(DataLoadingOption {
+ option_name: key.value,
+ option_type: DataLoadingOptionType::ENUM,
+ value: word.value,
+ }),
+ Token::Number(n, _) => Ok(DataLoadingOption {
+ option_name: key.value,
+ option_type: DataLoadingOptionType::NUMBER,
+ value: n,
+ }),
+ _ => parser.expected("expected option value", parser.peek_token()),
+ }
+ }
+}
+
/// Parsing a property of identity or autoincrement column option
/// Syntax:
/// ```sql
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
index c68ada3c..a18f1a4d 100644
--- a/tests/sqlparser_snowflake.rs
+++ b/tests/sqlparser_snowflake.rs
@@ -2028,20 +2028,25 @@ fn test_copy_into() {
);
match snowflake().verified_stmt(sql) {
Statement::CopyIntoSnowflake {
+ kind,
into,
- from_stage,
+ from_obj,
files,
pattern,
validation_mode,
..
} => {
+ assert_eq!(kind, CopyIntoSnowflakeKind::Table);
assert_eq!(
into,
ObjectName::from(vec![Ident::new("my_company"),
Ident::new("emp_basic")])
);
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::with_quote('\'',
"gcs://mybucket/./../a.csv")])
+ from_obj,
+ Some(ObjectName::from(vec![Ident::with_quote(
+ '\'',
+ "gcs://mybucket/./../a.csv"
+ )]))
);
assert!(files.is_none());
assert!(pattern.is_none());
@@ -2050,6 +2055,60 @@ fn test_copy_into() {
_ => unreachable!(),
};
assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
+
+ let sql = concat!("COPY INTO 's3://a/b/c/data.parquet' ", "FROM db.sc.tbl
", "PARTITION BY ('date=' || to_varchar(dt, 'YYYY-MM-DD') || '/hour=' ||
to_varchar(date_part(hour, ts)))");
+ match snowflake().verified_stmt(sql) {
+ Statement::CopyIntoSnowflake {
+ kind,
+ into,
+ from_obj,
+ from_query,
+ partition,
+ ..
+ } => {
+ assert_eq!(kind, CopyIntoSnowflakeKind::Location);
+ assert_eq!(
+ into,
+ ObjectName::from(vec![Ident::with_quote('\'',
"s3://a/b/c/data.parquet")])
+ );
+ assert_eq!(
+ from_obj,
+ Some(ObjectName::from(vec![
+ Ident::new("db"),
+ Ident::new("sc"),
+ Ident::new("tbl")
+ ]))
+ );
+ assert!(from_query.is_none());
+ assert!(partition.is_some());
+ }
+ _ => unreachable!(),
+ };
+ assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
+
+ let sql = concat!(
+ "COPY INTO 's3://a/b/c/data.parquet' ",
+ "FROM (SELECT * FROM tbl)"
+ );
+ match snowflake().verified_stmt(sql) {
+ Statement::CopyIntoSnowflake {
+ kind,
+ into,
+ from_obj,
+ from_query,
+ ..
+ } => {
+ assert_eq!(kind, CopyIntoSnowflakeKind::Location);
+ assert_eq!(
+ into,
+ ObjectName::from(vec![Ident::with_quote('\'',
"s3://a/b/c/data.parquet")])
+ );
+ assert!(from_query.is_some());
+ assert!(from_obj.is_none());
+ }
+ _ => unreachable!(),
+ };
+ assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
}
#[test]
@@ -2065,14 +2124,17 @@ fn test_copy_into_with_stage_params() {
match snowflake().verified_stmt(sql) {
Statement::CopyIntoSnowflake {
- from_stage,
+ from_obj,
stage_params,
..
} => {
//assert_eq!("s3://load/files/", stage_params.url.unwrap());
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::with_quote('\'',
"s3://load/files/")])
+ from_obj,
+ Some(ObjectName::from(vec![Ident::with_quote(
+ '\'',
+ "s3://load/files/"
+ )]))
);
assert_eq!("myint", stage_params.storage_integration.unwrap());
assert_eq!(
@@ -2125,13 +2187,16 @@ fn test_copy_into_with_stage_params() {
match snowflake().verified_stmt(sql) {
Statement::CopyIntoSnowflake {
- from_stage,
+ from_obj,
stage_params,
..
} => {
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::with_quote('\'',
"s3://load/files/")])
+ from_obj,
+ Some(ObjectName::from(vec![Ident::with_quote(
+ '\'',
+ "s3://load/files/"
+ )]))
);
assert_eq!("myint", stage_params.storage_integration.unwrap());
}
@@ -2154,13 +2219,13 @@ fn
test_copy_into_with_files_and_pattern_and_verification() {
files,
pattern,
validation_mode,
- from_stage_alias,
+ from_obj_alias,
..
} => {
assert_eq!(files.unwrap(), vec!["file1.json", "file2.json"]);
assert_eq!(pattern.unwrap(), ".*employees0[1-5].csv.gz");
assert_eq!(validation_mode.unwrap(), "RETURN_7_ROWS");
- assert_eq!(from_stage_alias.unwrap(), Ident::new("some_alias"));
+ assert_eq!(from_obj_alias.unwrap(), Ident::new("some_alias"));
}
_ => unreachable!(),
}
@@ -2179,13 +2244,16 @@ fn test_copy_into_with_transformations() {
match snowflake().verified_stmt(sql) {
Statement::CopyIntoSnowflake {
- from_stage,
+ from_obj,
from_transformations,
..
} => {
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::new("@schema"),
Ident::new("general_finished")])
+ from_obj,
+ Some(ObjectName::from(vec![
+ Ident::new("@schema"),
+ Ident::new("general_finished")
+ ]))
);
assert_eq!(
from_transformations.as_ref().unwrap()[0],
@@ -2254,6 +2322,41 @@ fn test_copy_into_file_format() {
snowflake_without_unescape().verified_stmt(sql).to_string(),
sql
);
+
+ // Test commas in file format
+ let sql = concat!(
+ "COPY INTO my_company.emp_basic ",
+ "FROM 'gcs://mybucket/./../a.csv' ",
+ "FILES = ('file1.json', 'file2.json') ",
+ "PATTERN = '.*employees0[1-5].csv.gz' ",
+ r#"FILE_FORMAT=(COMPRESSION=AUTO, BINARY_FORMAT=HEX, ESCAPE='\\')"#
+ );
+
+ match snowflake_without_unescape()
+ .parse_sql_statements(sql)
+ .unwrap()
+ .first()
+ .unwrap()
+ {
+ Statement::CopyIntoSnowflake { file_format, .. } => {
+ assert!(file_format.options.contains(&DataLoadingOption {
+ option_name: "COMPRESSION".to_string(),
+ option_type: DataLoadingOptionType::ENUM,
+ value: "AUTO".to_string()
+ }));
+ assert!(file_format.options.contains(&DataLoadingOption {
+ option_name: "BINARY_FORMAT".to_string(),
+ option_type: DataLoadingOptionType::ENUM,
+ value: "HEX".to_string()
+ }));
+ assert!(file_format.options.contains(&DataLoadingOption {
+ option_name: "ESCAPE".to_string(),
+ option_type: DataLoadingOptionType::STRING,
+ value: r#"\\"#.to_string()
+ }));
+ }
+ _ => unreachable!(),
+ }
}
#[test]
@@ -2285,16 +2388,8 @@ fn test_copy_into_copy_options() {
}
#[test]
-fn test_snowflake_stage_object_names() {
- let allowed_formatted_names = [
- "my_company.emp_basic",
- "@namespace.%table_name",
- "@namespace.%table_name/path",
- "@namespace.stage_name/path",
- "@~/path",
- ];
+fn test_snowflake_stage_object_names_into_location() {
let mut allowed_object_names = [
- ObjectName::from(vec![Ident::new("my_company"),
Ident::new("emp_basic")]),
ObjectName::from(vec![Ident::new("@namespace"),
Ident::new("%table_name")]),
ObjectName::from(vec![
Ident::new("@namespace"),
@@ -2307,7 +2402,39 @@ fn test_snowflake_stage_object_names() {
ObjectName::from(vec![Ident::new("@~/path")]),
];
- for it in allowed_formatted_names
+ let allowed_names_into_location = [
+ "@namespace.%table_name",
+ "@namespace.%table_name/path",
+ "@namespace.stage_name/path",
+ "@~/path",
+ ];
+ for it in allowed_names_into_location
+ .iter()
+ .zip(allowed_object_names.iter_mut())
+ {
+ let (formatted_name, object_name) = it;
+ let sql = format!(
+ "COPY INTO {} FROM 'gcs://mybucket/./../a.csv'",
+ formatted_name
+ );
+ match snowflake().verified_stmt(&sql) {
+ Statement::CopyIntoSnowflake { into, .. } => {
+ assert_eq!(into.0, object_name.0)
+ }
+ _ => unreachable!(),
+ }
+ }
+}
+
+#[test]
+fn test_snowflake_stage_object_names_into_table() {
+ let mut allowed_object_names = [
+ ObjectName::from(vec![Ident::new("my_company"),
Ident::new("emp_basic")]),
+ ObjectName::from(vec![Ident::new("emp_basic")]),
+ ];
+
+ let allowed_names_into_table = ["my_company.emp_basic", "emp_basic"];
+ for it in allowed_names_into_table
.iter()
.zip(allowed_object_names.iter_mut())
{
@@ -2330,16 +2457,17 @@ fn test_snowflake_copy_into() {
let sql = "COPY INTO a.b FROM @namespace.stage_name";
assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
match snowflake().verified_stmt(sql) {
- Statement::CopyIntoSnowflake {
- into, from_stage, ..
- } => {
+ Statement::CopyIntoSnowflake { into, from_obj, .. } => {
assert_eq!(
into,
ObjectName::from(vec![Ident::new("a"), Ident::new("b")])
);
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::new("@namespace"),
Ident::new("stage_name")])
+ from_obj,
+ Some(ObjectName::from(vec![
+ Ident::new("@namespace"),
+ Ident::new("stage_name")
+ ]))
)
}
_ => unreachable!(),
@@ -2351,9 +2479,7 @@ fn test_snowflake_copy_into_stage_name_ends_with_parens()
{
let sql = "COPY INTO SCHEMA.SOME_MONITORING_SYSTEM FROM (SELECT t.$1:st AS
st FROM @schema.general_finished)";
assert_eq!(snowflake().verified_stmt(sql).to_string(), sql);
match snowflake().verified_stmt(sql) {
- Statement::CopyIntoSnowflake {
- into, from_stage, ..
- } => {
+ Statement::CopyIntoSnowflake { into, from_obj, .. } => {
assert_eq!(
into,
ObjectName::from(vec![
@@ -2362,8 +2488,11 @@ fn
test_snowflake_copy_into_stage_name_ends_with_parens() {
])
);
assert_eq!(
- from_stage,
- ObjectName::from(vec![Ident::new("@schema"),
Ident::new("general_finished")])
+ from_obj,
+ Some(ObjectName::from(vec![
+ Ident::new("@schema"),
+ Ident::new("general_finished")
+ ]))
)
}
_ => unreachable!(),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
