set_tblproperties (#1151)

apache · Mar 1, 2024 · a511c47 · a511c47
1 parent 10cc54e
commit a511c47
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 4 deletions.
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
@@ -143,6 +143,8 @@ pub enum AlterTableOperation {
     ///
     /// Note: this is Snowflake specific <https://docs.snowflake.com/en/sql-reference/sql/alter-table>
     SwapWith { table_name: ObjectName },
+    /// 'SET TBLPROPERTIES ( { property_key [ = ] property_val } [, ...] )'
+    SetTblProperties { table_properties: Vec<SqlOption> },
 }
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
@@ -283,6 +285,13 @@ impl fmt::Display for AlterTableOperation {
             AlterTableOperation::SwapWith { table_name } => {
                 write!(f, "SWAP WITH {table_name}")
             }
+            AlterTableOperation::SetTblProperties { table_properties } => {
+                write!(
+                    f,
+                    "SET TBLPROPERTIES({})",
+                    display_comma_separated(table_properties)
+                )
+            }
         }
     }
 }

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -5124,6 +5124,20 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_options_with_keywords(
+        &mut self,
+        keywords: &[Keyword],
+    ) -> Result<Vec<SqlOption>, ParserError> {
+        if self.parse_keywords(keywords) {
+            self.expect_token(&Token::LParen)?;
+            let options = self.parse_comma_separated(Parser::parse_sql_option)?;
+            self.expect_token(&Token::RParen)?;
+            Ok(options)
+        } else {
+            Ok(vec![])
+        }
+    }
+
     pub fn parse_index_type(&mut self) -> Result<IndexType, ParserError> {
         if self.parse_keyword(Keyword::BTREE) {
             Ok(IndexType::BTree)
@@ -5385,10 +5399,18 @@ impl<'a> Parser<'a> {
             let table_name = self.parse_object_name(false)?;
             AlterTableOperation::SwapWith { table_name }
         } else {
-            return self.expected(
-                "ADD, RENAME, PARTITION, SWAP or DROP after ALTER TABLE",
-                self.peek_token(),
-            );
+            let options: Vec<SqlOption> =
+                self.parse_options_with_keywords(&[Keyword::SET, Keyword::TBLPROPERTIES])?;
+            if !options.is_empty() {
+                AlterTableOperation::SetTblProperties {
+                    table_properties: options,
+                }
+            } else {
+                return self.expected(
+                    "ADD, RENAME, PARTITION, SWAP, DROP, or SET TBLPROPERTIES after ALTER TABLE",
+                    self.peek_token(),
+                );
+            }
         };
         Ok(operation)
     }

diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -3539,6 +3539,23 @@ fn parse_alter_table() {
         }
         _ => unreachable!(),
     }
+
+    let set_table_properties = "ALTER TABLE tab SET TBLPROPERTIES('classification' = 'parquet')";
+    match alter_table_op(verified_stmt(set_table_properties)) {
+        AlterTableOperation::SetTblProperties { table_properties } => {
+            assert_eq!(
+                table_properties,
+                [SqlOption {
+                    name: Ident {
+                        value: "classification".to_string(),
+                        quote_style: Some('\'')
+                    },
+                    value: Expr::Value(Value::SingleQuotedString("parquet".to_string())),
+                }],
+            );
+        }
+        _ => unreachable!(),
+    }
 }
 
 #[test]