#[cfg(not(feature = "std"))]
use crate::alloc::string::ToString;
use crate::ast::helpers::stmt_data_loading::{
DataLoadingOption, DataLoadingOptionType, DataLoadingOptions, StageLoadSelectItem,
StageParamsObject,
};
use crate::ast::{Ident, ObjectName, Statement};
use crate::dialect::Dialect;
use crate::keywords::Keyword;
use crate::parser::{Parser, ParserError};
use crate::tokenizer::Token;
#[cfg(not(feature = "std"))]
use alloc::string::String;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(not(feature = "std"))]
use alloc::{format, vec};
#[derive(Debug, Default)]
pub struct SnowflakeDialect;
impl Dialect for SnowflakeDialect {
fn is_identifier_start(&self, ch: char) -> bool {
ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch == '_'
}
fn is_identifier_part(&self, ch: char) -> bool {
ch.is_ascii_lowercase()
|| ch.is_ascii_uppercase()
|| ch.is_ascii_digit()
|| ch == '$'
|| ch == '_'
}
fn supports_string_literal_backslash_escape(&self) -> bool {
true
}
fn supports_within_after_array_aggregation(&self) -> bool {
true
}
fn supports_connect_by(&self) -> bool {
true
}
fn supports_match_recognize(&self) -> bool {
true
}
fn supports_dictionary_syntax(&self) -> bool {
true
}
fn supports_window_function_null_treatment_arg(&self) -> bool {
true
}
fn supports_parenthesized_set_variables(&self) -> bool {
true
}
fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
if parser.parse_keyword(Keyword::CREATE) {
let or_replace = parser.parse_keywords(&[Keyword::OR, Keyword::REPLACE]);
let temporary = parser.parse_keyword(Keyword::TEMPORARY);
if parser.parse_keyword(Keyword::STAGE) {
return Some(parse_create_stage(or_replace, temporary, parser));
} else {
let mut back = 1;
if or_replace {
back += 2
}
if temporary {
back += 1
}
for _i in 0..back {
parser.prev_token();
}
}
}
if parser.parse_keywords(&[Keyword::COPY, Keyword::INTO]) {
return Some(parse_copy_into(parser));
}
None
}
}
pub fn parse_create_stage(
or_replace: bool,
temporary: bool,
parser: &mut Parser,
) -> Result<Statement, ParserError> {
let if_not_exists = parser.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
let name = parser.parse_object_name(false)?;
let mut directory_table_params = Vec::new();
let mut file_format = Vec::new();
let mut copy_options = Vec::new();
let mut comment = None;
let stage_params = parse_stage_params(parser)?;
if parser.parse_keyword(Keyword::DIRECTORY) {
parser.expect_token(&Token::Eq)?;
directory_table_params = parse_parentheses_options(parser)?;
}
if parser.parse_keyword(Keyword::FILE_FORMAT) {
parser.expect_token(&Token::Eq)?;
file_format = parse_parentheses_options(parser)?;
}
if parser.parse_keyword(Keyword::COPY_OPTIONS) {
parser.expect_token(&Token::Eq)?;
copy_options = parse_parentheses_options(parser)?;
}
if parser.parse_keyword(Keyword::COMMENT) {
parser.expect_token(&Token::Eq)?;
comment = Some(match parser.next_token().token {
Token::SingleQuotedString(word) => Ok(word),
_ => parser.expected("a comment statement", parser.peek_token()),
}?)
}
Ok(Statement::CreateStage {
or_replace,
temporary,
if_not_exists,
name,
stage_params,
directory_table_params: DataLoadingOptions {
options: directory_table_params,
},
file_format: DataLoadingOptions {
options: file_format,
},
copy_options: DataLoadingOptions {
options: copy_options,
},
comment,
})
}
pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserError> {
let mut ident = String::new();
while let Some(next_token) = parser.next_token_no_skip() {
match &next_token.token {
Token::Whitespace(_) => break,
Token::Period => {
parser.prev_token();
break;
}
Token::RParen => {
parser.prev_token();
break;
}
Token::AtSign => ident.push('@'),
Token::Tilde => ident.push('~'),
Token::Mod => ident.push('%'),
Token::Div => ident.push('/'),
Token::Word(w) => ident.push_str(&w.value),
_ => return parser.expected("stage name identifier", parser.peek_token()),
}
}
Ok(Ident::new(ident))
}
pub fn parse_snowflake_stage_name(parser: &mut Parser) -> Result<ObjectName, ParserError> {
match parser.next_token().token {
Token::AtSign => {
parser.prev_token();
let mut idents = vec![];
loop {
idents.push(parse_stage_name_identifier(parser)?);
if !parser.consume_token(&Token::Period) {
break;
}
}
Ok(ObjectName(idents))
}
_ => {
parser.prev_token();
Ok(parser.parse_object_name(false)?)
}
}
}
pub fn parse_copy_into(parser: &mut Parser) -> Result<Statement, ParserError> {
let into: ObjectName = parse_snowflake_stage_name(parser)?;
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;
parser.expect_keyword(Keyword::FROM)?;
match parser.next_token().token {
Token::LParen => {
parser.expect_keyword(Keyword::SELECT)?;
from_transformations = parse_select_items_for_data_load(parser)?;
parser.expect_keyword(Keyword::FROM)?;
from_stage = parse_snowflake_stage_name(parser)?;
stage_params = parse_stage_params(parser)?;
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
};
parser.expect_token(&Token::RParen)?;
}
_ => {
parser.prev_token();
from_stage = parse_snowflake_stage_name(parser)?;
stage_params = parse_stage_params(parser)?;
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
};
}
};
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(); }
}
parser.expect_token(&Token::RParen)?;
}
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)?,
});
}
let mut file_format = Vec::new();
if parser.parse_keyword(Keyword::FILE_FORMAT) {
parser.expect_token(&Token::Eq)?;
file_format = parse_parentheses_options(parser)?;
}
let mut copy_options = Vec::new();
if parser.parse_keyword(Keyword::COPY_OPTIONS) {
parser.expect_token(&Token::Eq)?;
copy_options = parse_parentheses_options(parser)?;
}
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 {
into,
from_stage,
from_stage_alias,
stage_params,
from_transformations,
files: if files.is_empty() { None } else { Some(files) },
pattern,
file_format: DataLoadingOptions {
options: file_format,
},
copy_options: DataLoadingOptions {
options: copy_options,
},
validation_mode,
})
}
fn parse_select_items_for_data_load(
parser: &mut Parser,
) -> Result<Option<Vec<StageLoadSelectItem>>, ParserError> {
let mut select_items: Vec<StageLoadSelectItem> = 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;
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)?;
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),
}?;
}
match parser.next_token().token {
Token::Colon => {
element = Some(Ident::new(match parser.next_token().token {
Token::Word(w) => Ok(w.value),
_ => parser.expected("file_col_num", parser.peek_token()),
}?));
}
_ => {
parser.prev_token();
}
}
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()),
}?);
}
select_items.push(StageLoadSelectItem {
alias,
file_col_num,
element,
item_as,
});
match parser.next_token().token {
Token::Comma => {
}
_ => {
parser.prev_token(); break;
}
}
}
Ok(Some(select_items))
}
fn parse_stage_params(parser: &mut Parser) -> Result<StageParamsObject, ParserError> {
let (mut url, mut storage_integration, mut endpoint) = (None, None, None);
let mut encryption: DataLoadingOptions = DataLoadingOptions { options: vec![] };
let mut credentials: DataLoadingOptions = DataLoadingOptions { options: vec![] };
if parser.parse_keyword(Keyword::URL) {
parser.expect_token(&Token::Eq)?;
url = Some(match parser.next_token().token {
Token::SingleQuotedString(word) => Ok(word),
_ => parser.expected("a URL statement", parser.peek_token()),
}?)
}
if parser.parse_keyword(Keyword::STORAGE_INTEGRATION) {
parser.expect_token(&Token::Eq)?;
storage_integration = Some(parser.next_token().token.to_string());
}
if parser.parse_keyword(Keyword::ENDPOINT) {
parser.expect_token(&Token::Eq)?;
endpoint = Some(match parser.next_token().token {
Token::SingleQuotedString(word) => Ok(word),
_ => parser.expected("an endpoint statement", parser.peek_token()),
}?)
}
if parser.parse_keyword(Keyword::CREDENTIALS) {
parser.expect_token(&Token::Eq)?;
credentials = DataLoadingOptions {
options: parse_parentheses_options(parser)?,
};
}
if parser.parse_keyword(Keyword::ENCRYPTION) {
parser.expect_token(&Token::Eq)?;
encryption = DataLoadingOptions {
options: parse_parentheses_options(parser)?,
};
}
Ok(StageParamsObject {
url,
encryption,
endpoint,
storage_integration,
credentials,
})
}
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()),
}?;
}
Ok(options)
}