Added support for DROP OPERATOR FAMILY (apache#2106)

Author: LucaCappelletti94

src/ast/ddl.rs

@@ -4251,3 +4251,40 @@ impl Spanned for DropOperator {
         Span::empty()
     }
 }
+
+/// `DROP OPERATOR FAMILY` statement
+/// See <https://www.postgresql.org/docs/current/sql-dropopfamily.html>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct DropOperatorFamily {
+    /// `IF EXISTS` clause
+    pub if_exists: bool,
+    /// One or more operator families to drop
+    pub names: Vec<ObjectName>,
+    /// Index method (btree, hash, gist, gin, etc.)
+    pub using: Ident,
+    /// `CASCADE or RESTRICT`
+    pub drop_behavior: Option<DropBehavior>,
+}
+
+impl fmt::Display for DropOperatorFamily {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "DROP OPERATOR FAMILY")?;
+        if self.if_exists {
+            write!(f, " IF EXISTS")?;
+        }
+        write!(f, " {}", display_comma_separated(&self.names))?;
+        write!(f, " USING {}", self.using)?;
+        if let Some(drop_behavior) = &self.drop_behavior {
+            write!(f, " {}", drop_behavior)?;
+        }
+        Ok(())
+    }
+}
+
+impl Spanned for DropOperatorFamily {
+    fn span(&self) -> Span {
+        Span::empty()
+    }
+}

src/ast/mod.rs

@@ -67,13 +67,13 @@ pub use self::ddl::{
     ColumnPolicyProperty, ConstraintCharacteristics, CreateConnector, CreateDomain,
     CreateExtension, CreateFunction, CreateIndex, CreateOperator, CreateOperatorClass,
     CreateOperatorFamily, CreateTable, CreateTrigger, CreateView, Deduplicate, DeferrableInitial,
-    DropBehavior, DropExtension, DropFunction, DropOperator, DropOperatorSignature, DropTrigger,
-    GeneratedAs, GeneratedExpressionMode, IdentityParameters, IdentityProperty,
-    IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder, IndexColumn,
-    IndexOption, IndexType, KeyOrIndexDisplay, Msck, NullsDistinctOption, OperatorArgTypes,
-    OperatorClassItem, OperatorPurpose, Owner, Partition, ProcedureParam, ReferentialAction,
-    RenameTableNameKind, ReplicaIdentity, TagsColumnOption, TriggerObjectKind, Truncate,
-    UserDefinedTypeCompositeAttributeDef, UserDefinedTypeInternalLength,
+    DropBehavior, DropExtension, DropFunction, DropOperator, DropOperatorFamily,
+    DropOperatorSignature, DropTrigger, GeneratedAs, GeneratedExpressionMode, IdentityParameters,
+    IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder,
+    IndexColumn, IndexOption, IndexType, KeyOrIndexDisplay, Msck, NullsDistinctOption,
+    OperatorArgTypes, OperatorClassItem, OperatorPurpose, Owner, Partition, ProcedureParam,
+    ReferentialAction, RenameTableNameKind, ReplicaIdentity, TagsColumnOption, TriggerObjectKind,
+    Truncate, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeInternalLength,
     UserDefinedTypeRangeOption, UserDefinedTypeRepresentation, UserDefinedTypeSqlDefinitionOption,
     UserDefinedTypeStorage, ViewColumnDef,
 };
@@ -3580,6 +3580,12 @@ pub enum Statement {
     /// <https://www.postgresql.org/docs/current/sql-dropoperator.html>
     DropOperator(DropOperator),
     /// ```sql
+    /// DROP OPERATOR FAMILY [ IF EXISTS ] name USING index_method [ CASCADE | RESTRICT ]
+    /// ```
+    /// Note: this is a PostgreSQL-specific statement.
+    /// <https://www.postgresql.org/docs/current/sql-dropopfamily.html>
+    DropOperatorFamily(DropOperatorFamily),
+    /// ```sql
     /// FETCH
     /// ```
     /// Retrieve rows from a query using a cursor
@@ -4844,6 +4850,9 @@ impl fmt::Display for Statement {
             Statement::CreateExtension(create_extension) => write!(f, "{create_extension}"),
             Statement::DropExtension(drop_extension) => write!(f, "{drop_extension}"),
             Statement::DropOperator(drop_operator) => write!(f, "{drop_operator}"),
+            Statement::DropOperatorFamily(drop_operator_family) => {
+                write!(f, "{drop_operator_family}")
+            }
             Statement::CreateRole(create_role) => write!(f, "{create_role}"),
             Statement::CreateSecret {
                 or_replace,

src/ast/spans.rs

@@ -376,6 +376,7 @@ impl Spanned for Statement {
             Statement::CreateExtension(create_extension) => create_extension.span(),
             Statement::DropExtension(drop_extension) => drop_extension.span(),
             Statement::DropOperator(drop_operator) => drop_operator.span(),
+            Statement::DropOperatorFamily(drop_operator_family) => drop_operator_family.span(),
             Statement::CreateSecret { .. } => Span::empty(),
             Statement::CreateServer { .. } => Span::empty(),
             Statement::CreateConnector { .. } => Span::empty(),

src/parser/mod.rs

@@ -6773,7 +6773,12 @@ impl<'a> Parser<'a> {
         } else if self.parse_keyword(Keyword::EXTENSION) {
             return self.parse_drop_extension();
         } else if self.parse_keyword(Keyword::OPERATOR) {
-            return self.parse_drop_operator();
+            // Check if this is DROP OPERATOR FAMILY
+            return if self.parse_keyword(Keyword::FAMILY) {
+                self.parse_drop_operator_family()
+            } else {
+                self.parse_drop_operator()
+            };
         } else {
             return self.expected(
                 "CONNECTOR, DATABASE, EXTENSION, FUNCTION, INDEX, OPERATOR, POLICY, PROCEDURE, ROLE, SCHEMA, SECRET, SEQUENCE, STAGE, TABLE, TRIGGER, TYPE, VIEW, MATERIALIZED VIEW or USER after DROP",
@@ -7572,6 +7577,23 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a [Statement::DropOperatorFamily]
+    ///
+    /// [PostgreSQL Documentation](https://www.postgresql.org/docs/current/sql-dropopfamily.html)
+    pub fn parse_drop_operator_family(&mut self) -> Result<Statement, ParserError> {
+        let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+        let names = self.parse_comma_separated(|p| p.parse_object_name(false))?;
+        self.expect_keyword(Keyword::USING)?;
+        let using = self.parse_identifier()?;
+        let drop_behavior = self.parse_optional_drop_behavior();
+        Ok(Statement::DropOperatorFamily(DropOperatorFamily {
+            if_exists,
+            names,
+            using,
+            drop_behavior,
+        }))
+    }
+
     //TODO: Implement parsing for Skewed
     pub fn parse_hive_distribution(&mut self) -> Result<HiveDistributionStyle, ParserError> {
         if self.parse_keywords(&[Keyword::PARTITIONED, Keyword::BY]) {

tests/sqlparser_postgres.rs

@@ -6910,6 +6910,61 @@ fn parse_drop_operator() {
     );
 }
 
+#[test]
+fn parse_drop_operator_family() {
+    for if_exists in [true, false] {
+        for drop_behavior in [
+            None,
+            Some(DropBehavior::Cascade),
+            Some(DropBehavior::Restrict),
+        ] {
+            for index_method in &["btree", "hash", "gist", "gin", "spgist", "brin"] {
+                for (names_str, names_vec) in [
+                    (
+                        "float_ops",
+                        vec![ObjectName::from(vec![Ident::new("float_ops")])],
+                    ),
+                    (
+                        "myschema.custom_ops",
+                        vec![ObjectName::from(vec![
+                            Ident::new("myschema"),
+                            Ident::new("custom_ops"),
+                        ])],
+                    ),
+                    (
+                        "ops1, ops2, schema.ops3",
+                        vec![
+                            ObjectName::from(vec![Ident::new("ops1")]),
+                            ObjectName::from(vec![Ident::new("ops2")]),
+                            ObjectName::from(vec![Ident::new("schema"), Ident::new("ops3")]),
+                        ],
+                    ),
+                ] {
+                    let sql = format!(
+                        "DROP OPERATOR FAMILY{} {} USING {}{}",
+                        if if_exists { " IF EXISTS" } else { "" },
+                        names_str,
+                        index_method,
+                        match drop_behavior {
+                            Some(behavior) => format!(" {}", behavior),
+                            None => String::new(),
+                        }
+                    );
+                    assert_eq!(
+                        pg_and_generic().verified_stmt(&sql),
+                        Statement::DropOperatorFamily(DropOperatorFamily {
+                            if_exists,
+                            names: names_vec,
+                            using: Ident::new(*index_method),
+                            drop_behavior,
+                        })
+                    );
+                }
+            }
+        }
+    }
+}
+
 #[test]
 fn parse_create_operator_family() {
     for index_method in &["btree", "hash", "gist", "gin", "spgist", "brin"] {