Support Snowflake ALTER STAGE syntax

src/ast/helpers/stmt_data_loading.rs

@@ -48,6 +48,43 @@ pub struct StageParamsObject {
     pub credentials: KeyValueOptions,
 }
 
+impl StageParamsObject {
+    /// Returns true when no stage parameter is present.
+    pub fn is_empty(&self) -> bool {
+        self.url.is_none()
+            && self.encryption.options.is_empty()
+            && self.endpoint.is_none()
+            && self.storage_integration.is_none()
+            && self.credentials.options.is_empty()
+    }
+}
+
+/// An operation supported by Snowflake's `ALTER STAGE` statement.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AlterStageOperation {
+    /// Rename a stage.
+    RenameTo {
+        /// New stage name.
+        #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
+        new_name: ObjectName,
+    },
+    /// Set one or more stage properties.
+    Set {
+        /// External stage parameters.
+        stage_params: StageParamsObject,
+        /// Directory table parameters.
+        directory_table_params: KeyValueOptions,
+        /// File format options.
+        file_format: KeyValueOptions,
+        /// Copy options.
+        copy_options: KeyValueOptions,
+        /// Optional comment.
+        comment: Option<String>,
+    },
+}
+
 /// This enum enables support for both standard SQL select item expressions
 /// and Snowflake-specific ones for data loading.
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]

src/ast/mod.rs

@@ -26,7 +26,7 @@ use alloc::{
 };
 use helpers::{
     attached_token::AttachedToken,
-    stmt_data_loading::{FileStagingCommand, StageLoadSelectItemKind},
+    stmt_data_loading::{AlterStageOperation, FileStagingCommand, StageLoadSelectItemKind},
 };
 
 use core::cmp::Ordering;
@@ -4535,6 +4535,18 @@ pub enum Statement {
         comment: Option<String>,
     },
     /// ```sql
+    /// ALTER STAGE
+    /// ```
+    /// See <https://docs.snowflake.com/en/sql-reference/sql/alter-stage>
+    AlterStage {
+        /// `IF EXISTS` flag.
+        if_exists: bool,
+        /// Stage name.
+        name: ObjectName,
+        /// Stage alteration to perform.
+        operation: AlterStageOperation,
+    },
+    /// ```sql
     /// CREATE FILE FORMAT
     /// ```
     /// See <https://docs.snowflake.com/en/sql-reference/sql/create-file-format>
@@ -6306,6 +6318,44 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::AlterStage {
+                if_exists,
+                name,
+                operation,
+            } => {
+                write!(
+                    f,
+                    "ALTER STAGE {if_exists}{name}",
+                    if_exists = if *if_exists { "IF EXISTS " } else { "" },
+                )?;
+                match operation {
+                    AlterStageOperation::RenameTo { new_name } => {
+                        write!(f, " RENAME TO {new_name}")
+                    }
+                    AlterStageOperation::Set {
+                        stage_params,
+                        directory_table_params,
+                        file_format,
+                        copy_options,
+                        comment,
+                    } => {
+                        write!(f, " SET{stage_params}")?;
+                        if !directory_table_params.options.is_empty() {
+                            write!(f, " DIRECTORY=({directory_table_params})")?;
+                        }
+                        if !file_format.options.is_empty() {
+                            write!(f, " FILE_FORMAT=({file_format})")?;
+                        }
+                        if !copy_options.options.is_empty() {
+                            write!(f, " COPY_OPTIONS=({copy_options})")?;
+                        }
+                        if let Some(comment) = comment {
+                            write!(f, " COMMENT='{comment}'")?;
+                        }
+                        Ok(())
+                    }
+                }
+            }
             Statement::CreateFileFormat {
                 or_replace,
                 temporary,

src/ast/spans.rs

@@ -299,6 +299,7 @@ impl Spanned for Values {
 /// - [Statement::CreateProcedure]
 /// - [Statement::CreateMacro]
 /// - [Statement::CreateStage]
+/// - [Statement::AlterStage]
 /// - [Statement::Assert]
 /// - [Statement::Grant]
 /// - [Statement::Revoke]
@@ -463,6 +464,7 @@ impl Spanned for Statement {
             Statement::CreateProcedure { .. } => Span::empty(),
             Statement::CreateMacro { .. } => Span::empty(),
             Statement::CreateStage { .. } => Span::empty(),
+            Statement::AlterStage { .. } => Span::empty(),
             Statement::CreateFileFormat { .. } => Span::empty(),
             Statement::Assert { .. } => Span::empty(),
             Statement::Grant { .. } => Span::empty(),

src/dialect/snowflake.rs

@@ -24,7 +24,8 @@ use crate::ast::helpers::key_value_options::{
 use crate::ast::helpers::stmt_create_database::CreateDatabaseBuilder;
 use crate::ast::helpers::stmt_create_table::CreateTableBuilder;
 use crate::ast::helpers::stmt_data_loading::{
-    FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
+    AlterStageOperation, FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind,
+    StageParamsObject,
 };
 use crate::ast::{
     AlterTable, AlterTableOperation, AlterTableType, CatalogSyncNamespaceMode, CloudProviderParams,
@@ -298,6 +299,10 @@ impl Dialect for SnowflakeDialect {
             return Some(parse_alter_session(parser, set));
         }
 
+        if parser.parse_keywords(&[Keyword::ALTER, Keyword::STAGE]) {
+            return Some(parse_alter_stage(parser));
+        }
+
         if parser.parse_keyword(Keyword::CREATE) {
             // possibly CREATE STAGE
             //[ OR  REPLACE ]
@@ -1397,60 +1402,63 @@ pub fn parse_create_stage(
     //[ IF NOT EXISTS ]
     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;
 
     // [ internalStageParams | externalStageParams ]
     let stage_params = parse_stage_params(parser)?;
-
-    // [ directoryTableParams ]
-    if parser.parse_keyword(Keyword::DIRECTORY) {
-        parser.expect_token(&Token::Eq)?;
-        directory_table_params = parser.parse_key_value_options(true, &[])?.options;
-    }
-
-    // [ file_format]
-    if parser.parse_keyword(Keyword::FILE_FORMAT) {
-        parser.expect_token(&Token::Eq)?;
-        file_format = parser.parse_key_value_options(true, &[])?.options;
-    }
-
-    // [ copy_options ]
-    if parser.parse_keyword(Keyword::COPY_OPTIONS) {
-        parser.expect_token(&Token::Eq)?;
-        copy_options = parser.parse_key_value_options(true, &[])?.options;
-    }
-
-    // [ comment ]
-    if parser.parse_keyword(Keyword::COMMENT) {
-        parser.expect_token(&Token::Eq)?;
-        comment = Some(parser.parse_comment_value()?);
-    }
+    let (directory_table_params, file_format, copy_options, comment) = parse_stage_options(parser)?;
 
     Ok(Statement::CreateStage {
         or_replace,
         temporary,
         if_not_exists,
         name,
         stage_params,
-        directory_table_params: KeyValueOptions {
-            options: directory_table_params,
-            delimiter: KeyValueOptionsDelimiter::Space,
-        },
-        file_format: KeyValueOptions {
-            options: file_format,
-            delimiter: KeyValueOptionsDelimiter::Space,
-        },
-        copy_options: KeyValueOptions {
-            options: copy_options,
-            delimiter: KeyValueOptionsDelimiter::Space,
-        },
+        directory_table_params,
+        file_format,
+        copy_options,
         comment,
     })
 }
 
+fn parse_alter_stage(parser: &mut Parser) -> Result<Statement, ParserError> {
+    let if_exists = parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+    let name = parser.parse_object_name(false)?;
+
+    let operation = if parser.parse_keyword(Keyword::RENAME) {
+        parser.expect_keyword(Keyword::TO)?;
+        AlterStageOperation::RenameTo {
+            new_name: parser.parse_object_name(false)?,
+        }
+    } else if parser.parse_keyword(Keyword::SET) {
+        let stage_params = parse_stage_params(parser)?;
+        let (directory_table_params, file_format, copy_options, comment) =
+            parse_stage_options(parser)?;
+        if stage_params.is_empty()
+            && directory_table_params.options.is_empty()
+            && file_format.options.is_empty()
+            && copy_options.options.is_empty()
+            && comment.is_none()
+        {
+            return parser.expected_ref("a stage property", parser.peek_token_ref());
+        }
+        AlterStageOperation::Set {
+            stage_params,
+            directory_table_params,
+            file_format,
+            copy_options,
+            comment,
+        }
+    } else {
+        return parser.expected_ref("RENAME TO or SET", parser.peek_token_ref());
+    };
+
+    Ok(Statement::AlterStage {
+        if_exists,
+        name,
+        operation,
+    })
+}
+
 pub fn parse_create_file_format(
     or_replace: bool,
     temporary: bool,
@@ -1858,6 +1866,56 @@ fn parse_stage_params(parser: &mut Parser) -> Result<StageParamsObject, ParserEr
     })
 }
 
+fn parse_stage_options(
+    parser: &mut Parser,
+) -> Result<
+    (
+        KeyValueOptions,
+        KeyValueOptions,
+        KeyValueOptions,
+        Option<String>,
+    ),
+    ParserError,
+> {
+    let mut directory_table_params = Vec::new();
+    let mut file_format = Vec::new();
+    let mut copy_options = Vec::new();
+    let mut comment = None;
+
+    if parser.parse_keyword(Keyword::DIRECTORY) {
+        parser.expect_token(&Token::Eq)?;
+        directory_table_params = parser.parse_key_value_options(true, &[])?.options;
+    }
+    if parser.parse_keyword(Keyword::FILE_FORMAT) {
+        parser.expect_token(&Token::Eq)?;
+        file_format = parser.parse_key_value_options(true, &[])?.options;
+    }
+    if parser.parse_keyword(Keyword::COPY_OPTIONS) {
+        parser.expect_token(&Token::Eq)?;
+        copy_options = parser.parse_key_value_options(true, &[])?.options;
+    }
+    if parser.parse_keyword(Keyword::COMMENT) {
+        parser.expect_token(&Token::Eq)?;
+        comment = Some(parser.parse_comment_value()?);
+    }
+
+    Ok((
+        KeyValueOptions {
+            options: directory_table_params,
+            delimiter: KeyValueOptionsDelimiter::Space,
+        },
+        KeyValueOptions {
+            options: file_format,
+            delimiter: KeyValueOptionsDelimiter::Space,
+        },
+        KeyValueOptions {
+            options: copy_options,
+            delimiter: KeyValueOptionsDelimiter::Space,
+        },
+        comment,
+    ))
+}
+
 /// Parses options separated by blank spaces, commas, or new lines like:
 /// ABORT_DETACHED_QUERY = { TRUE | FALSE }
 ///      [ ACTIVE_PYTHON_PROFILER = { 'LINE' | 'MEMORY' } ]

tests/sqlparser_snowflake.rs

@@ -20,7 +20,9 @@
 //! generic dialect is also tested (on the inputs it can handle).
 
 use sqlparser::ast::helpers::key_value_options::{KeyValueOption, KeyValueOptionKind};
-use sqlparser::ast::helpers::stmt_data_loading::{StageLoadSelectItem, StageLoadSelectItemKind};
+use sqlparser::ast::helpers::stmt_data_loading::{
+    AlterStageOperation, StageLoadSelectItem, StageLoadSelectItemKind,
+};
 use sqlparser::ast::*;
 use sqlparser::dialect::{Dialect, GenericDialect, SnowflakeDialect};
 use sqlparser::parser::{ParserError, ParserOptions};
@@ -2185,6 +2187,72 @@ fn test_create_stage() {
     );
 }
 
+#[test]
+fn test_alter_stage() {
+    let rename_sql = "ALTER STAGE IF EXISTS analytics.raw.events RENAME TO archived_events";
+    match snowflake().verified_stmt(rename_sql) {
+        Statement::AlterStage {
+            if_exists,
+            name,
+            operation: AlterStageOperation::RenameTo { new_name },
+        } => {
+            assert!(if_exists);
+            assert_eq!("analytics.raw.events", name.to_string());
+            assert_eq!("archived_events", new_name.to_string());
+        }
+        _ => unreachable!(),
+    }
+    assert_eq!(
+        snowflake().verified_stmt(rename_sql).to_string(),
+        rename_sql
+    );
+
+    let set_sql = concat!(
+        "ALTER STAGE analytics.raw.events SET ",
+        "URL='s3://bucket/path/' STORAGE_INTEGRATION=my_int ",
+        "ENDPOINT='s3.us-east-2.amazonaws.com' ",
+        "CREDENTIALS=(AWS_KEY_ID='key' AWS_SECRET_KEY='secret') ",
+        "ENCRYPTION=(TYPE='AWS_SSE_KMS' KMS_KEY_ID='id') ",
+        "DIRECTORY=(ENABLE=true AUTO_REFRESH=false) ",
+        "FILE_FORMAT=(TYPE=PARQUET BINARY_AS_TEXT=false) ",
+        "COPY_OPTIONS=(ON_ERROR='SKIP_FILE') COMMENT='updated'"
+    );
+    match snowflake().verified_stmt(set_sql) {
+        Statement::AlterStage {
+            if_exists,
+            name,
+            operation:
+                AlterStageOperation::Set {
+                    stage_params,
+                    directory_table_params,
+                    file_format,
+                    copy_options,
+                    comment,
+                },
+        } => {
+            assert!(!if_exists);
+            assert_eq!("analytics.raw.events", name.to_string());
+            assert_eq!(Some("s3://bucket/path/"), stage_params.url.as_deref());
+            assert_eq!(Some("my_int"), stage_params.storage_integration.as_deref());
+            assert_eq!(2, stage_params.credentials.options.len());
+            assert_eq!(2, stage_params.encryption.options.len());
+            assert_eq!(2, directory_table_params.options.len());
+            assert_eq!(2, file_format.options.len());
+            assert_eq!(1, copy_options.options.len());
+            assert_eq!(Some("updated"), comment.as_deref());
+        }
+        _ => unreachable!(),
+    }
+    assert_eq!(snowflake().verified_stmt(set_sql).to_string(), set_sql);
+
+    for unsupported in [
+        "ALTER STAGE analytics.raw.events SET",
+        "ALTER STAGE analytics.raw.events REFRESH",
+    ] {
+        assert!(snowflake().parse_sql_statements(unsupported).is_err());
+    }
+}
+
 #[test]
 fn test_create_file_format() {
     let sql = "CREATE FILE FORMAT analytics.formats.parquet TYPE=PARQUET";