Skip to content

Commit

Permalink
Snowflake ALTER TABLE clustering options (#1579)
Browse files Browse the repository at this point in the history
  • Loading branch information
yoavcloud authored Dec 6, 2024
1 parent 7b50ac3 commit d0fcc06
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 15 deletions.
83 changes: 68 additions & 15 deletions src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,10 @@ pub enum AlterTableOperation {
///
/// Note: this is a ClickHouse-specific operation.
/// Please refer to [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/alter/projection#drop-projection)
DropProjection { if_exists: bool, name: Ident },
DropProjection {
if_exists: bool,
name: Ident,
},

/// `MATERIALIZE PROJECTION [IF EXISTS] name [IN PARTITION partition_name]`
///
Expand Down Expand Up @@ -99,11 +102,15 @@ pub enum AlterTableOperation {
/// `DISABLE RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
DisableRule { name: Ident },
DisableRule {
name: Ident,
},
/// `DISABLE TRIGGER [ trigger_name | ALL | USER ]`
///
/// Note: this is a PostgreSQL-specific operation.
DisableTrigger { name: Ident },
DisableTrigger {
name: Ident,
},
/// `DROP CONSTRAINT [ IF EXISTS ] <name>`
DropConstraint {
if_exists: bool,
Expand Down Expand Up @@ -152,31 +159,43 @@ pub enum AlterTableOperation {
/// `ENABLE ALWAYS RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableAlwaysRule { name: Ident },
EnableAlwaysRule {
name: Ident,
},
/// `ENABLE ALWAYS TRIGGER trigger_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableAlwaysTrigger { name: Ident },
EnableAlwaysTrigger {
name: Ident,
},
/// `ENABLE REPLICA RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableReplicaRule { name: Ident },
EnableReplicaRule {
name: Ident,
},
/// `ENABLE REPLICA TRIGGER trigger_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableReplicaTrigger { name: Ident },
EnableReplicaTrigger {
name: Ident,
},
/// `ENABLE ROW LEVEL SECURITY`
///
/// Note: this is a PostgreSQL-specific operation.
EnableRowLevelSecurity,
/// `ENABLE RULE rewrite_rule_name`
///
/// Note: this is a PostgreSQL-specific operation.
EnableRule { name: Ident },
EnableRule {
name: Ident,
},
/// `ENABLE TRIGGER [ trigger_name | ALL | USER ]`
///
/// Note: this is a PostgreSQL-specific operation.
EnableTrigger { name: Ident },
EnableTrigger {
name: Ident,
},
/// `RENAME TO PARTITION (partition=val)`
RenamePartitions {
old_partitions: Vec<Expr>,
Expand All @@ -197,7 +216,9 @@ pub enum AlterTableOperation {
new_column_name: Ident,
},
/// `RENAME TO <table_name>`
RenameTable { table_name: ObjectName },
RenameTable {
table_name: ObjectName,
},
// CHANGE [ COLUMN ] <old_name> <new_name> <data_type> [ <options> ]
ChangeColumn {
old_name: Ident,
Expand All @@ -218,7 +239,10 @@ pub enum AlterTableOperation {
/// `RENAME CONSTRAINT <old_constraint_name> TO <new_constraint_name>`
///
/// Note: this is a PostgreSQL-specific operation.
RenameConstraint { old_name: Ident, new_name: Ident },
RenameConstraint {
old_name: Ident,
new_name: Ident,
},
/// `ALTER [ COLUMN ]`
AlterColumn {
column_name: Ident,
Expand All @@ -227,14 +251,27 @@ pub enum AlterTableOperation {
/// 'SWAP WITH <table_name>'
///
/// Note: this is Snowflake specific <https://docs.snowflake.com/en/sql-reference/sql/alter-table>
SwapWith { table_name: ObjectName },
SwapWith {
table_name: ObjectName,
},
/// 'SET TBLPROPERTIES ( { property_key [ = ] property_val } [, ...] )'
SetTblProperties { table_properties: Vec<SqlOption> },

SetTblProperties {
table_properties: Vec<SqlOption>,
},
/// `OWNER TO { <new_owner> | CURRENT_ROLE | CURRENT_USER | SESSION_USER }`
///
/// Note: this is PostgreSQL-specific <https://www.postgresql.org/docs/current/sql-altertable.html>
OwnerTo { new_owner: Owner },
OwnerTo {
new_owner: Owner,
},
/// Snowflake table clustering options
/// <https://docs.snowflake.com/en/sql-reference/sql/alter-table#clustering-actions-clusteringaction>
ClusterBy {
exprs: Vec<Expr>,
},
DropClusteringKey,
SuspendRecluster,
ResumeRecluster,
}

/// An `ALTER Policy` (`Statement::AlterPolicy`) operation
Expand Down Expand Up @@ -548,6 +585,22 @@ impl fmt::Display for AlterTableOperation {
}
Ok(())
}
AlterTableOperation::ClusterBy { exprs } => {
write!(f, "CLUSTER BY ({})", display_comma_separated(exprs))?;
Ok(())
}
AlterTableOperation::DropClusteringKey => {
write!(f, "DROP CLUSTERING KEY")?;
Ok(())
}
AlterTableOperation::SuspendRecluster => {
write!(f, "SUSPEND RECLUSTER")?;
Ok(())
}
AlterTableOperation::ResumeRecluster => {
write!(f, "RESUME RECLUSTER")?;
Ok(())
}
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions src/ast/spans.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1020,6 +1020,10 @@ impl Spanned for AlterTableOperation {
union_spans(table_properties.iter().map(|i| i.span()))
}
AlterTableOperation::OwnerTo { .. } => Span::empty(),
AlterTableOperation::ClusterBy { exprs } => union_spans(exprs.iter().map(|e| e.span())),
AlterTableOperation::DropClusteringKey => Span::empty(),
AlterTableOperation::SuspendRecluster => Span::empty(),
AlterTableOperation::ResumeRecluster => Span::empty(),
}
}
}
Expand Down
4 changes: 4 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ define_keywords!(
CLOSE,
CLUSTER,
CLUSTERED,
CLUSTERING,
COALESCE,
COLLATE,
COLLATION,
Expand Down Expand Up @@ -622,6 +623,7 @@ define_keywords!(
READS,
READ_ONLY,
REAL,
RECLUSTER,
RECURSIVE,
REF,
REFERENCES,
Expand Down Expand Up @@ -656,6 +658,7 @@ define_keywords!(
RESTRICTIVE,
RESULT,
RESULTSET,
RESUME,
RETAIN,
RETURN,
RETURNING,
Expand Down Expand Up @@ -745,6 +748,7 @@ define_keywords!(
SUM,
SUPER,
SUPERUSER,
SUSPEND,
SWAP,
SYMMETRIC,
SYNC,
Expand Down
11 changes: 11 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7273,6 +7273,8 @@ impl<'a> Parser<'a> {
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
let name = self.parse_identifier(false)?;
AlterTableOperation::DropProjection { if_exists, name }
} else if self.parse_keywords(&[Keyword::CLUSTERING, Keyword::KEY]) {
AlterTableOperation::DropClusteringKey
} else {
let _ = self.parse_keyword(Keyword::COLUMN); // [ COLUMN ]
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
Expand Down Expand Up @@ -7444,6 +7446,15 @@ impl<'a> Parser<'a> {
partition,
with_name,
}
} else if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
self.expect_token(&Token::LParen)?;
let exprs = self.parse_comma_separated(|parser| parser.parse_expr())?;
self.expect_token(&Token::RParen)?;
AlterTableOperation::ClusterBy { exprs }
} else if self.parse_keywords(&[Keyword::SUSPEND, Keyword::RECLUSTER]) {
AlterTableOperation::SuspendRecluster
} else if self.parse_keywords(&[Keyword::RESUME, Keyword::RECLUSTER]) {
AlterTableOperation::ResumeRecluster
} else {
let options: Vec<SqlOption> =
self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;
Expand Down
36 changes: 36 additions & 0 deletions tests/sqlparser_snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1411,6 +1411,42 @@ fn test_alter_table_swap_with() {
};
}

#[test]
fn test_alter_table_clustering() {
let sql = r#"ALTER TABLE tab CLUSTER BY (c1, "c2", TO_DATE(c3))"#;
match alter_table_op(snowflake_and_generic().verified_stmt(sql)) {
AlterTableOperation::ClusterBy { exprs } => {
assert_eq!(
exprs,
[
Expr::Identifier(Ident::new("c1")),
Expr::Identifier(Ident::with_quote('"', "c2")),
Expr::Function(Function {
name: ObjectName(vec![Ident::new("TO_DATE")]),
parameters: FunctionArguments::None,
args: FunctionArguments::List(FunctionArgumentList {
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::Identifier(Ident::new("c3"))
))],
duplicate_treatment: None,
clauses: vec![],
}),
filter: None,
null_treatment: None,
over: None,
within_group: vec![]
})
],
);
}
_ => unreachable!(),
}

snowflake_and_generic().verified_stmt("ALTER TABLE tbl DROP CLUSTERING KEY");
snowflake_and_generic().verified_stmt("ALTER TABLE tbl SUSPEND RECLUSTER");
snowflake_and_generic().verified_stmt("ALTER TABLE tbl RESUME RECLUSTER");
}

#[test]
fn test_drop_stage() {
match snowflake_and_generic().verified_stmt("DROP STAGE s1") {
Expand Down

0 comments on commit d0fcc06

Please sign in to comment.