Add drop behavior to DROP PRIMARY/FOREIGN KEY

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>`
     ///
@@ -658,8 +659,31 @@ impl fmt::Display for AlterTableOperation {
                     }
                 )
             }
-            AlterTableOperation::DropPrimaryKey => write!(f, "DROP PRIMARY KEY"),
-            AlterTableOperation::DropForeignKey { name } => write!(f, "DROP FOREIGN KEY {name}"),
+            AlterTableOperation::DropPrimaryKey { drop_behavior } => {
+                write!(
+                    f,
+                    "DROP PRIMARY KEY{}",
+                    match drop_behavior {
+                        None => "",
+                        Some(DropBehavior::Restrict) => " RESTRICT",
+                        Some(DropBehavior::Cascade) => " CASCADE",
+                    }
+                )
+            }
+            AlterTableOperation::DropForeignKey {
+                name,
+                drop_behavior,
+            } => {
+                write!(
+                    f,
+                    "DROP FOREIGN KEY {name}{}",
+                    match drop_behavior {
+                        None => "",
+                        Some(DropBehavior::Restrict) => " RESTRICT",
+                        Some(DropBehavior::Cascade) => " CASCADE",
+                    }
+                )
+            }
             AlterTableOperation::DropIndex { name } => write!(f, "DROP INDEX {name}"),
             AlterTableOperation::DropColumn {
                 has_column_keyword,

src/ast/spans.rs

@@ -1151,8 +1151,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

@@ -8896,10 +8896,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");
+}