Add support for DROP EXTENSION (apache#1610)

goldmedal · Dec 27, 2024 · 7dbf31b · 7dbf31b
1 parent d89cf80
commit 7dbf31b
Show file tree

Hide file tree

Showing 4 changed files with 144 additions and 1 deletion.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -2755,6 +2755,18 @@ pub enum Statement {
         version: Option<Ident>,
     },
     /// ```sql
+    /// DROP EXTENSION [ IF EXISTS ] name [, ...] [ CASCADE | RESTRICT ]
+    ///
+    /// Note: this is a PostgreSQL-specific statement.
+    /// https://www.postgresql.org/docs/current/sql-dropextension.html
+    /// ```
+    DropExtension {
+        names: Vec<Ident>,
+        if_exists: bool,
+        /// `CASCADE` or `RESTRICT`
+        cascade_or_restrict: Option<ReferentialAction>,
+    },
+    /// ```sql
     /// FETCH
     /// ```
     /// Retrieve rows from a query using a cursor
@@ -4029,6 +4041,21 @@ impl fmt::Display for Statement {
 
                 Ok(())
             }
+            Statement::DropExtension {
+                names,
+                if_exists,
+                cascade_or_restrict,
+            } => {
+                write!(f, "DROP EXTENSION")?;
+                if *if_exists {
+                    write!(f, " IF EXISTS")?;
+                }
+                write!(f, " {}", display_comma_separated(names))?;
+                if let Some(cascade_or_restrict) = cascade_or_restrict {
+                    write!(f, " {cascade_or_restrict}")?;
+                }
+                Ok(())
+            }
             Statement::CreateRole {
                 names,
                 if_not_exists,

diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -430,6 +430,7 @@ impl Spanned for Statement {
             Statement::DropSecret { .. } => Span::empty(),
             Statement::Declare { .. } => Span::empty(),
             Statement::CreateExtension { .. } => Span::empty(),
+            Statement::DropExtension { .. } => Span::empty(),
             Statement::Fetch { .. } => Span::empty(),
             Statement::Flush { .. } => Span::empty(),
             Statement::Discard { .. } => Span::empty(),

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -5414,9 +5414,11 @@ impl<'a> Parser<'a> {
             return self.parse_drop_secret(temporary, persistent);
         } else if self.parse_keyword(Keyword::TRIGGER) {
             return self.parse_drop_trigger();
+        } else if self.parse_keyword(Keyword::EXTENSION) {
+            return self.parse_drop_extension();
         } else {
             return self.expected(
-                "TABLE, VIEW, INDEX, ROLE, SCHEMA, DATABASE, FUNCTION, PROCEDURE, STAGE, TRIGGER, SECRET, SEQUENCE, or TYPE after DROP",
+                "TABLE, VIEW, INDEX, ROLE, SCHEMA, DATABASE, FUNCTION, PROCEDURE, STAGE, TRIGGER, SECRET, SEQUENCE, TYPE, or EXTENSION after DROP",
                 self.peek_token(),
             );
         };
@@ -6079,6 +6081,25 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a PostgreSQL-specific [Statement::DropExtension] statement.
+    pub fn parse_drop_extension(&mut self) -> Result<Statement, ParserError> {
+        let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+        let names = self.parse_comma_separated(|p| p.parse_identifier(false))?;
+        let cascade_or_restrict =
+            self.parse_one_of_keywords(&[Keyword::CASCADE, Keyword::RESTRICT]);
+        Ok(Statement::DropExtension {
+            names,
+            if_exists,
+            cascade_or_restrict: cascade_or_restrict
+                .map(|k| match k {
+                    Keyword::CASCADE => Ok(ReferentialAction::Cascade),
+                    Keyword::RESTRICT => Ok(ReferentialAction::Restrict),
+                    _ => self.expected("CASCADE or RESTRICT", self.peek_token()),
+                })
+                .transpose()?,
+        })
+    }
+
     //TODO: Implement parsing for Skewed
     pub fn parse_hive_distribution(&mut self) -> Result<HiveDistributionStyle, ParserError> {
         if self.parse_keywords(&[Keyword::PARTITIONED, Keyword::BY]) {

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -662,6 +662,100 @@ fn parse_create_extension() {
         .verified_stmt("CREATE EXTENSION extension_name WITH SCHEMA schema_name VERSION version");
 }
 
+#[test]
+fn parse_drop_extension() {
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION extension_name"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: false,
+            cascade_or_restrict: None,
+        }
+    );
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION extension_name CASCADE"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: false,
+            cascade_or_restrict: Some(ReferentialAction::Cascade),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION extension_name RESTRICT"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: false,
+            cascade_or_restrict: Some(ReferentialAction::Restrict),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION extension_name, extension_name2 CASCADE"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into(), "extension_name2".into()],
+            if_exists: false,
+            cascade_or_restrict: Some(ReferentialAction::Cascade),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION extension_name, extension_name2 RESTRICT"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into(), "extension_name2".into()],
+            if_exists: false,
+            cascade_or_restrict: Some(ReferentialAction::Restrict),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION IF EXISTS extension_name"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: true,
+            cascade_or_restrict: None,
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION IF EXISTS extension_name CASCADE"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: true,
+            cascade_or_restrict: Some(ReferentialAction::Cascade),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic().verified_stmt("DROP EXTENSION IF EXISTS extension_name RESTRICT"),
+        Statement::DropExtension {
+            names: vec!["extension_name".into()],
+            if_exists: true,
+            cascade_or_restrict: Some(ReferentialAction::Restrict),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic()
+            .verified_stmt("DROP EXTENSION IF EXISTS extension_name1, extension_name2 CASCADE"),
+        Statement::DropExtension {
+            names: vec!["extension_name1".into(), "extension_name2".into()],
+            if_exists: true,
+            cascade_or_restrict: Some(ReferentialAction::Cascade),
+        }
+    );
+
+    assert_eq!(
+        pg_and_generic()
+            .verified_stmt("DROP EXTENSION IF EXISTS extension_name1, extension_name2 RESTRICT"),
+        Statement::DropExtension {
+            names: vec!["extension_name1".into(), "extension_name2".into()],
+            if_exists: true,
+            cascade_or_restrict: Some(ReferentialAction::Restrict),
+        }
+    );
+}
+
 #[test]
 fn parse_alter_table_alter_column() {
     pg().one_statement_parses_to(