diff --git a/src/ast/helpers/stmt_data_loading.rs b/src/ast/helpers/stmt_data_loading.rs index cda6c6ea4..42e1df06b 100644 --- a/src/ast/helpers/stmt_data_loading.rs +++ b/src/ast/helpers/stmt_data_loading.rs @@ -29,7 +29,7 @@ use core::fmt::Formatter; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; -use crate::ast::Ident; +use crate::ast::{Ident, ObjectName}; #[cfg(feature = "visitor")] use sqlparser_derive::{Visit, VisitMut}; @@ -156,3 +156,22 @@ impl fmt::Display for StageLoadSelectItem { Ok(()) } } + +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub struct FileStagingCommand { + #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] + pub stage: ObjectName, + pub pattern: Option, +} + +impl fmt::Display for FileStagingCommand { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.stage)?; + if let Some(pattern) = self.pattern.as_ref() { + write!(f, " PATTERN='{pattern}'")?; + } + Ok(()) + } +} diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 867ae25b6..f46438b3e 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -23,7 +23,7 @@ use alloc::{ string::{String, ToString}, vec::Vec, }; -use helpers::attached_token::AttachedToken; +use helpers::{attached_token::AttachedToken, stmt_data_loading::FileStagingCommand}; use core::ops::Deref; use core::{ @@ -3420,6 +3420,12 @@ pub enum Statement { /// /// See Mysql RenameTable(Vec), + /// Snowflake `LIST` + /// See: + List(FileStagingCommand), + /// Snowflake `REMOVE` + /// See: + Remove(FileStagingCommand), } impl fmt::Display for Statement { @@ -4980,6 +4986,8 @@ impl fmt::Display for Statement { Statement::RenameTable(rename_tables) => { write!(f, "RENAME TABLE {}", display_comma_separated(rename_tables)) } + Statement::List(command) => write!(f, "LIST {command}"), + Statement::Remove(command) => write!(f, "REMOVE {command}"), } } } diff --git a/src/ast/spans.rs b/src/ast/spans.rs index dad0c5379..1dd9118f5 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -493,6 +493,7 @@ impl Spanned for Statement { Statement::LoadData { .. } => Span::empty(), Statement::UNLISTEN { .. } => Span::empty(), Statement::RenameTable { .. } => Span::empty(), + Statement::List(..) | Statement::Remove(..) => Span::empty(), } } } diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs index 249241d73..55343da18 100644 --- a/src/dialect/snowflake.rs +++ b/src/dialect/snowflake.rs @@ -19,8 +19,8 @@ use crate::alloc::string::ToString; use crate::ast::helpers::stmt_create_table::CreateTableBuilder; use crate::ast::helpers::stmt_data_loading::{ - DataLoadingOption, DataLoadingOptionType, DataLoadingOptions, StageLoadSelectItem, - StageParamsObject, + DataLoadingOption, DataLoadingOptionType, DataLoadingOptions, FileStagingCommand, + StageLoadSelectItem, StageParamsObject, }; use crate::ast::{ ColumnOption, ColumnPolicy, ColumnPolicyProperty, Ident, IdentityParameters, IdentityProperty, @@ -165,6 +165,15 @@ impl Dialect for SnowflakeDialect { return Some(parse_copy_into(parser)); } + if let Some(kw) = parser.parse_one_of_keywords(&[ + Keyword::LIST, + Keyword::LS, + Keyword::REMOVE, + Keyword::RM, + ]) { + return Some(parse_file_staging_command(kw, parser)); + } + None } @@ -240,6 +249,26 @@ impl Dialect for SnowflakeDialect { } } +fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result { + let stage = parse_snowflake_stage_name(parser)?; + let pattern = if parser.parse_keyword(Keyword::PATTERN) { + parser.expect_token(&Token::Eq)?; + Some(parser.parse_literal_string()?) + } else { + None + }; + + match kw { + Keyword::LIST | Keyword::LS => Ok(Statement::List(FileStagingCommand { stage, pattern })), + Keyword::REMOVE | Keyword::RM => { + Ok(Statement::Remove(FileStagingCommand { stage, pattern })) + } + _ => Err(ParserError::ParserError( + "unexpected stage command, expecting LIST, LS, REMOVE or RM".to_string(), + )), + } +} + /// Parse snowflake create table statement. /// pub fn parse_create_table( @@ -501,7 +530,7 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result ident.push('~'), Token::Mod => ident.push('%'), Token::Div => ident.push('/'), - Token::Word(w) => ident.push_str(&w.value), + Token::Word(w) => ident.push_str(&w.to_string()), _ => return parser.expected("stage name identifier", parser.peek_token()), } } diff --git a/src/keywords.rs b/src/keywords.rs index 3fed882c3..b7ff39e04 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -451,6 +451,7 @@ define_keywords!( LIKE_REGEX, LIMIT, LINES, + LIST, LISTEN, LN, LOAD, @@ -467,6 +468,7 @@ define_keywords!( LOWCARDINALITY, LOWER, LOW_PRIORITY, + LS, MACRO, MANAGEDLOCATION, MAP, @@ -649,6 +651,7 @@ define_keywords!( RELAY, RELEASE, REMOTE, + REMOVE, RENAME, REORG, REPAIR, @@ -672,6 +675,7 @@ define_keywords!( REVOKE, RIGHT, RLIKE, + RM, ROLE, ROLES, ROLLBACK, diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index b5add6116..112aa5264 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -2991,3 +2991,34 @@ fn test_table_sample() { snowflake_and_generic().verified_stmt("SELECT id FROM mytable TABLESAMPLE (10) REPEATABLE (1)"); snowflake_and_generic().verified_stmt("SELECT id FROM mytable TABLESAMPLE (10) SEED (1)"); } + +#[test] +fn parse_ls_and_rm() { + snowflake().one_statement_parses_to("LS @~", "LIST @~"); + snowflake().one_statement_parses_to("RM @~", "REMOVE @~"); + + let statement = snowflake() + .verified_stmt("LIST @SNOWFLAKE_KAFKA_CONNECTOR_externalDataLakeSnowflakeConnector_STAGE_call_tracker_stream/"); + match statement { + Statement::List(command) => { + assert_eq!(command.stage, ObjectName(vec!["@SNOWFLAKE_KAFKA_CONNECTOR_externalDataLakeSnowflakeConnector_STAGE_call_tracker_stream/".into()])); + assert!(command.pattern.is_none()); + } + _ => unreachable!(), + }; + + let statement = + snowflake().verified_stmt("REMOVE @my_csv_stage/analysis/ PATTERN='.*data_0.*'"); + match statement { + Statement::Remove(command) => { + assert_eq!( + command.stage, + ObjectName(vec!["@my_csv_stage/analysis/".into()]) + ); + assert_eq!(command.pattern, Some(".*data_0.*".to_string())); + } + _ => unreachable!(), + }; + + snowflake().verified_stmt(r#"LIST @"STAGE_WITH_QUOTES""#); +}