Add drop behavior to DROP PRIMARY/FOREIGN KEY (apache#2002)

Author: yoavcloud

src/ast/ddl.rs

@@ -200,17 +200,18 @@ pub enum AlterTableOperation {
     },
     /// `DROP PRIMARY KEY`
     ///
-    /// Note: this is a [MySQL]-specific operation.
-    ///
-    /// [MySQL]: https://dev.mysql.com/doc/refman/8.4/en/alter-table.html
-    DropPrimaryKey,
+    /// [MySQL](https://dev.mysql.com/doc/refman/8.4/en/alter-table.html)
+    /// [Snowflake](https://docs.snowflake.com/en/sql-reference/constraints-drop)
+    DropPrimaryKey {
+        drop_behavior: Option<DropBehavior>,
+    },
     /// `DROP FOREIGN KEY <fk_symbol>`
     ///
-    /// Note: this is a [MySQL]-specific operation.
-    ///
-    /// [MySQL]: https://dev.mysql.com/doc/refman/8.4/en/alter-table.html
+    /// [MySQL](https://dev.mysql.com/doc/refman/8.4/en/alter-table.html)
+    /// [Snowflake](https://docs.snowflake.com/en/sql-reference/constraints-drop)
     DropForeignKey {
         name: Ident,
+        drop_behavior: Option<DropBehavior>,
     },
     /// `DROP INDEX <index_name>`
     ///
@@ -648,36 +649,51 @@ impl fmt::Display for AlterTableOperation {
             } => {
                 write!(
                     f,
-                    "DROP CONSTRAINT {}{}{}",
+                    "DROP CONSTRAINT {}{}",
                     if *if_exists { "IF EXISTS " } else { "" },
-                    name,
-                    match drop_behavior {
-                        None => "",
-                        Some(DropBehavior::Restrict) => " RESTRICT",
-                        Some(DropBehavior::Cascade) => " CASCADE",
-                    }
-                )
+                    name
+                )?;
+                if let Some(drop_behavior) = drop_behavior {
+                    write!(f, " {drop_behavior}")?;
+                }
+                Ok(())
+            }
+            AlterTableOperation::DropPrimaryKey { drop_behavior } => {
+                write!(f, "DROP PRIMARY KEY")?;
+                if let Some(drop_behavior) = drop_behavior {
+                    write!(f, " {drop_behavior}")?;
+                }
+                Ok(())
+            }
+            AlterTableOperation::DropForeignKey {
+                name,
+                drop_behavior,
+            } => {
+                write!(f, "DROP FOREIGN KEY {name}")?;
+                if let Some(drop_behavior) = drop_behavior {
+                    write!(f, " {drop_behavior}")?;
+                }
+                Ok(())
             }
-            AlterTableOperation::DropPrimaryKey => write!(f, "DROP PRIMARY KEY"),
-            AlterTableOperation::DropForeignKey { name } => write!(f, "DROP FOREIGN KEY {name}"),
             AlterTableOperation::DropIndex { name } => write!(f, "DROP INDEX {name}"),
             AlterTableOperation::DropColumn {
                 has_column_keyword,
                 column_names: column_name,
                 if_exists,
                 drop_behavior,
-            } => write!(
-                f,
-                "DROP {}{}{}{}",
-                if *has_column_keyword { "COLUMN " } else { "" },
-                if *if_exists { "IF EXISTS " } else { "" },
-                display_comma_separated(column_name),
-                match drop_behavior {
-                    None => "",
-                    Some(DropBehavior::Restrict) => " RESTRICT",
-                    Some(DropBehavior::Cascade) => " CASCADE",
+            } => {
+                write!(
+                    f,
+                    "DROP {}{}{}",
+                    if *has_column_keyword { "COLUMN " } else { "" },
+                    if *if_exists { "IF EXISTS " } else { "" },
+                    display_comma_separated(column_name),
+                )?;
+                if let Some(drop_behavior) = drop_behavior {
+                    write!(f, " {drop_behavior}")?;
                 }
-            ),
+                Ok(())
+            }
             AlterTableOperation::AttachPartition { partition } => {
                 write!(f, "ATTACH {partition}")
             }

src/ast/spans.rs

@@ -1150,8 +1150,8 @@ impl Spanned for AlterTableOperation {
             } => partition
                 .span()
                 .union_opt(&with_name.as_ref().map(|n| n.span)),
-            AlterTableOperation::DropPrimaryKey => Span::empty(),
-            AlterTableOperation::DropForeignKey { name } => name.span,
+            AlterTableOperation::DropPrimaryKey { .. } => Span::empty(),
+            AlterTableOperation::DropForeignKey { name, .. } => name.span,
             AlterTableOperation::DropIndex { name } => name.span,
             AlterTableOperation::EnableAlwaysRule { name } => name.span,
             AlterTableOperation::EnableAlwaysTrigger { name } => name.span,

src/parser/mod.rs

@@ -8930,10 +8930,15 @@ impl<'a> Parser<'a> {
                     drop_behavior,
                 }
             } else if self.parse_keywords(&[Keyword::PRIMARY, Keyword::KEY]) {
-                AlterTableOperation::DropPrimaryKey
+                let drop_behavior = self.parse_optional_drop_behavior();
+                AlterTableOperation::DropPrimaryKey { drop_behavior }
             } else if self.parse_keywords(&[Keyword::FOREIGN, Keyword::KEY]) {
                 let name = self.parse_identifier()?;
-                AlterTableOperation::DropForeignKey { name }
+                let drop_behavior = self.parse_optional_drop_behavior();
+                AlterTableOperation::DropForeignKey {
+                    name,
+                    drop_behavior,
+                }
             } else if self.parse_keyword(Keyword::INDEX) {
                 let name = self.parse_identifier()?;
                 AlterTableOperation::DropIndex { name }

tests/sqlparser_mysql.rs

@@ -2752,7 +2752,9 @@ fn parse_alter_table_add_columns() {
 fn parse_alter_table_drop_primary_key() {
     assert_matches!(
         alter_table_op(mysql_and_generic().verified_stmt("ALTER TABLE tab DROP PRIMARY KEY")),
-        AlterTableOperation::DropPrimaryKey
+        AlterTableOperation::DropPrimaryKey {
+            drop_behavior: None
+        }
     );
 }
 
@@ -2762,7 +2764,7 @@ fn parse_alter_table_drop_foreign_key() {
         alter_table_op(
             mysql_and_generic().verified_stmt("ALTER TABLE tab DROP FOREIGN KEY foo_ibfk_1")
         ),
-        AlterTableOperation::DropForeignKey { name } if name.value == "foo_ibfk_1"
+        AlterTableOperation::DropForeignKey { name, .. } if name.value == "foo_ibfk_1"
     );
 }
 

tests/sqlparser_snowflake.rs

@@ -4571,3 +4571,13 @@ fn test_create_database() {
         .to_string();
     assert!(err.contains("Expected"), "Unexpected error: {err}");
 }
+
+#[test]
+fn test_drop_constraints() {
+    snowflake().verified_stmt("ALTER TABLE tbl DROP PRIMARY KEY");
+    snowflake().verified_stmt("ALTER TABLE tbl DROP FOREIGN KEY k1");
+    snowflake().verified_stmt("ALTER TABLE tbl DROP CONSTRAINT c1");
+    snowflake().verified_stmt("ALTER TABLE tbl DROP PRIMARY KEY CASCADE");
+    snowflake().verified_stmt("ALTER TABLE tbl DROP FOREIGN KEY k1 RESTRICT");
+    snowflake().verified_stmt("ALTER TABLE tbl DROP CONSTRAINT c1 CASCADE");
+}