Add support for Snowflake LIST and REMOVE (#1639)

Author: yoavcloud

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<String>,
+}
+
+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(())
+    }
+}

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 <https://dev.mysql.com/doc/refman/9.1/en/rename-table.html>
     RenameTable(Vec<RenameTable>),
+    /// Snowflake `LIST`
+    /// See: <https://docs.snowflake.com/en/sql-reference/sql/list>
+    List(FileStagingCommand),
+    /// Snowflake `REMOVE`
+    /// See: <https://docs.snowflake.com/en/sql-reference/sql/remove>
+    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}"),
         }
     }
 }

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(),
         }
     }
 }

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<Statement, ParserError> {
+    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.
 /// <https://docs.snowflake.com/en/sql-reference/sql/create-table>
 pub fn parse_create_table(
@@ -501,7 +530,7 @@ pub fn parse_stage_name_identifier(parser: &mut Parser) -> Result<Ident, ParserE
             Token::Tilde => 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()),
         }
     }

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,

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""#);
+}