PostgreSQL: Share operator-purpose parsing

Author: LucaCappelletti94

src/parser/mod.rs

@@ -8040,20 +8040,7 @@ impl<'a> Parser<'a> {
                     None
                 };
 
-                // Optional purpose
-                let purpose = if self.parse_keyword(Keyword::FOR) {
-                    if self.parse_keyword(Keyword::SEARCH) {
-                        Some(OperatorPurpose::ForSearch)
-                    } else if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
-                        let sort_family = self.parse_object_name(false)?;
-                        Some(OperatorPurpose::ForOrderBy { sort_family })
-                    } else {
-                        return self
-                            .expected_ref("SEARCH or ORDER BY after FOR", self.peek_token_ref());
-                    }
-                } else {
-                    None
-                };
+                let purpose = self.parse_operator_purpose()?;
 
                 items.push(OperatorClassItem::Operator {
                     strategy_number,
@@ -12298,6 +12285,21 @@ impl<'a> Parser<'a> {
         })
     }
 
+    fn parse_operator_purpose(&mut self) -> Result<Option<OperatorPurpose>, ParserError> {
+        if !self.parse_keyword(Keyword::FOR) {
+            return Ok(None);
+        }
+
+        if self.parse_keyword(Keyword::SEARCH) {
+            Ok(Some(OperatorPurpose::ForSearch))
+        } else if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
+            let sort_family = self.parse_object_name(false)?;
+            Ok(Some(OperatorPurpose::ForOrderBy { sort_family }))
+        } else {
+            self.expected_ref("SEARCH or ORDER BY after FOR", self.peek_token_ref())
+        }
+    }
+
     /// Parse an operator item for ALTER OPERATOR FAMILY ADD operations
     fn parse_operator_family_add_operator(&mut self) -> Result<OperatorFamilyItem, ParserError> {
         let strategy_number = self.parse_literal_uint()?;
@@ -12308,19 +12310,7 @@ impl<'a> Parser<'a> {
         let op_types = self.parse_comma_separated(Parser::parse_data_type)?;
         self.expect_token(&Token::RParen)?;
 
-        // Optional purpose
-        let purpose = if self.parse_keyword(Keyword::FOR) {
-            if self.parse_keyword(Keyword::SEARCH) {
-                Some(OperatorPurpose::ForSearch)
-            } else if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
-                let sort_family = self.parse_object_name(false)?;
-                Some(OperatorPurpose::ForOrderBy { sort_family })
-            } else {
-                return self.expected_ref("SEARCH or ORDER BY after FOR", self.peek_token_ref());
-            }
-        } else {
-            None
-        };
+        let purpose = self.parse_operator_purpose()?;
 
         Ok(OperatorFamilyItem::Operator {
             strategy_number,

tests/sqlparser_postgres.rs

@@ -11026,3 +11026,37 @@ fn parse_create_foreign_data_wrapper_clause_order() {
         pg_and_generic().verified_stmt(sql);
     }
 }
+
+#[test]
+fn parse_alter_operator_family_order_by_followed_by_search() {
+    let sql = r#"ALTER OPERATOR FAMILY my_ops USING gist ADD OPERATOR 1 < (INT4, INT4) FOR ORDER BY sort_schema."sort family", OPERATOR 2 > (INT4, INT4) FOR SEARCH"#;
+    let Statement::AlterOperatorFamily(AlterOperatorFamily {
+        operation: AlterOperatorFamilyOperation::Add { items },
+        ..
+    }) = pg_and_generic().verified_stmt(sql)
+    else {
+        panic!("Expected ALTER OPERATOR FAMILY");
+    };
+    assert_eq!(
+        items,
+        vec![
+            OperatorFamilyItem::Operator {
+                strategy_number: 1,
+                operator_name: ObjectName::from(vec![Ident::new("<")]),
+                op_types: vec![DataType::Int4(None), DataType::Int4(None)],
+                purpose: Some(OperatorPurpose::ForOrderBy {
+                    sort_family: ObjectName::from(vec![
+                        Ident::new("sort_schema"),
+                        Ident::with_quote('"', "sort family"),
+                    ]),
+                }),
+            },
+            OperatorFamilyItem::Operator {
+                strategy_number: 2,
+                operator_name: ObjectName::from(vec![Ident::new(">")]),
+                op_types: vec![DataType::Int4(None), DataType::Int4(None)],
+                purpose: Some(OperatorPurpose::ForSearch),
+            },
+        ]
+    );
+}