Add support for dropping multiple columns in Snowflake (#1918)

This commit is contained in:
Yoav Cohen 2025-07-05 08:18:58 +02:00 committed by GitHub
parent b0bcc46e22
commit d2466af20a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 28 additions and 12 deletions

View file

@ -140,10 +140,10 @@ pub enum AlterTableOperation {
name: Ident, name: Ident,
drop_behavior: Option<DropBehavior>, drop_behavior: Option<DropBehavior>,
}, },
/// `DROP [ COLUMN ] [ IF EXISTS ] <column_name> [ CASCADE ]` /// `DROP [ COLUMN ] [ IF EXISTS ] <column_name> [ , <column_name>, ... ] [ CASCADE ]`
DropColumn { DropColumn {
has_column_keyword: bool, has_column_keyword: bool,
column_name: Ident, column_names: Vec<Ident>,
if_exists: bool, if_exists: bool,
drop_behavior: Option<DropBehavior>, drop_behavior: Option<DropBehavior>,
}, },
@ -631,7 +631,7 @@ impl fmt::Display for AlterTableOperation {
AlterTableOperation::DropIndex { name } => write!(f, "DROP INDEX {name}"), AlterTableOperation::DropIndex { name } => write!(f, "DROP INDEX {name}"),
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword, has_column_keyword,
column_name, column_names: column_name,
if_exists, if_exists,
drop_behavior, drop_behavior,
} => write!( } => write!(
@ -639,7 +639,7 @@ impl fmt::Display for AlterTableOperation {
"DROP {}{}{}{}", "DROP {}{}{}{}",
if *has_column_keyword { "COLUMN " } else { "" }, if *has_column_keyword { "COLUMN " } else { "" },
if *if_exists { "IF EXISTS " } else { "" }, if *if_exists { "IF EXISTS " } else { "" },
column_name, display_comma_separated(column_name),
match drop_behavior { match drop_behavior {
None => "", None => "",
Some(DropBehavior::Restrict) => " RESTRICT", Some(DropBehavior::Restrict) => " RESTRICT",

View file

@ -1112,10 +1112,10 @@ impl Spanned for AlterTableOperation {
} => name.span, } => name.span,
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword: _, has_column_keyword: _,
column_name, column_names,
if_exists: _, if_exists: _,
drop_behavior: _, drop_behavior: _,
} => column_name.span, } => union_spans(column_names.iter().map(|i| i.span)),
AlterTableOperation::AttachPartition { partition } => partition.span(), AlterTableOperation::AttachPartition { partition } => partition.span(),
AlterTableOperation::DetachPartition { partition } => partition.span(), AlterTableOperation::DetachPartition { partition } => partition.span(),
AlterTableOperation::FreezePartition { AlterTableOperation::FreezePartition {

View file

@ -1071,6 +1071,11 @@ pub trait Dialect: Debug + Any {
fn supports_alter_column_type_using(&self) -> bool { fn supports_alter_column_type_using(&self) -> bool {
false false
} }
/// Returns true if the dialect supports `ALTER TABLE tbl DROP COLUMN c1, ..., cn`
fn supports_comma_separated_drop_column_list(&self) -> bool {
false
}
} }
/// This represents the operators for which precedence must be defined /// This represents the operators for which precedence must be defined

View file

@ -364,6 +364,10 @@ impl Dialect for SnowflakeDialect {
fn supports_space_separated_column_options(&self) -> bool { fn supports_space_separated_column_options(&self) -> bool {
true true
} }
fn supports_comma_separated_drop_column_list(&self) -> bool {
true
}
} }
fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> { fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {

View file

@ -8675,11 +8675,15 @@ impl<'a> Parser<'a> {
} else { } else {
let has_column_keyword = self.parse_keyword(Keyword::COLUMN); // [ COLUMN ] let has_column_keyword = self.parse_keyword(Keyword::COLUMN); // [ COLUMN ]
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]); let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
let column_name = self.parse_identifier()?; let column_names = if self.dialect.supports_comma_separated_drop_column_list() {
self.parse_comma_separated(Parser::parse_identifier)?
} else {
vec![self.parse_identifier()?]
};
let drop_behavior = self.parse_optional_drop_behavior(); let drop_behavior = self.parse_optional_drop_behavior();
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword, has_column_keyword,
column_name, column_names,
if_exists, if_exists,
drop_behavior, drop_behavior,
} }

View file

@ -4996,15 +4996,18 @@ fn parse_alter_table_drop_column() {
"ALTER TABLE tab DROP is_active CASCADE", "ALTER TABLE tab DROP is_active CASCADE",
); );
let dialects = all_dialects_where(|d| d.supports_comma_separated_drop_column_list());
dialects.verified_stmt("ALTER TABLE tbl DROP COLUMN c1, c2, c3");
fn check_one(constraint_text: &str) { fn check_one(constraint_text: &str) {
match alter_table_op(verified_stmt(&format!("ALTER TABLE tab {constraint_text}"))) { match alter_table_op(verified_stmt(&format!("ALTER TABLE tab {constraint_text}"))) {
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword: true, has_column_keyword: true,
column_name, column_names,
if_exists, if_exists,
drop_behavior, drop_behavior,
} => { } => {
assert_eq!("is_active", column_name.to_string()); assert_eq!("is_active", column_names.first().unwrap().to_string());
assert!(if_exists); assert!(if_exists);
match drop_behavior { match drop_behavior {
None => assert!(constraint_text.ends_with(" is_active")), None => assert!(constraint_text.ends_with(" is_active")),

View file

@ -2876,7 +2876,7 @@ fn parse_alter_table_with_algorithm() {
vec![ vec![
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword: true, has_column_keyword: true,
column_name: Ident::new("password_digest"), column_names: vec![Ident::new("password_digest")],
if_exists: false, if_exists: false,
drop_behavior: None, drop_behavior: None,
}, },
@ -2924,7 +2924,7 @@ fn parse_alter_table_with_lock() {
vec![ vec![
AlterTableOperation::DropColumn { AlterTableOperation::DropColumn {
has_column_keyword: true, has_column_keyword: true,
column_name: Ident::new("password_digest"), column_names: vec![Ident::new("password_digest")],
if_exists: false, if_exists: false,
drop_behavior: None, drop_behavior: None,
}, },