From 3e973a87585f0640b72e308a37c74f72b71e5724 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Mon, 21 Sep 2026 23:06:50 +0300 Subject: [PATCH] Support Snowflake ALTER STAGE syntax --- src/ast/helpers/stmt_data_loading.rs | 37 +++++++ src/ast/mod.rs | 52 +++++++++- src/ast/spans.rs | 2 + src/dialect/snowflake.rs | 140 +++++++++++++++++++-------- tests/sqlparser_snowflake.rs | 70 +++++++++++++- 5 files changed, 258 insertions(+), 43 deletions(-) diff --git a/src/ast/helpers/stmt_data_loading.rs b/src/ast/helpers/stmt_data_loading.rs index dfc1f4b0bf..56e91dd4ba 100644 --- a/src/ast/helpers/stmt_data_loading.rs +++ b/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, + }, +} + /// 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)] diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 79a4d7bcd6..15ca6a5cbc 100644 --- a/src/ast/mod.rs +++ b/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, }, /// ```sql + /// ALTER STAGE + /// ``` + /// See + AlterStage { + /// `IF EXISTS` flag. + if_exists: bool, + /// Stage name. + name: ObjectName, + /// Stage alteration to perform. + operation: AlterStageOperation, + }, + /// ```sql /// CREATE FILE FORMAT /// ``` /// See @@ -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, diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 446b87ae05..1223420a11 100644 --- a/src/ast/spans.rs +++ b/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(), diff --git a/src/dialect/snowflake.rs b/src/dialect/snowflake.rs index eeec285a57..4ec0f96550 100644 --- a/src/dialect/snowflake.rs +++ b/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,37 +1402,10 @@ 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, @@ -1435,22 +1413,52 @@ pub fn parse_create_stage( 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 { + 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 Result< + ( + KeyValueOptions, + KeyValueOptions, + KeyValueOptions, + Option, + ), + 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' } ] diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index d4bf5611e4..49d324220d 100644 --- a/tests/sqlparser_snowflake.rs +++ b/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";