[postgres] Add support for custom binary operators (apache#548)

* [postgres] Add support for custom binary operators

More details about operators in general are at:
https://www.postgresql.org/docs/current/sql-createoperator.html. This
patch attempts to parse `SELECT` queries that reference an operator
using `OPERATOR(<optional_schema>.<operator_name>)` syntax.

This is a PostgreSQL extension. There are no provisions for user-defined operators in the SQL standard.

* fix code-review comments and ci failures

* Allow custom operator in generic dialects too

* Parse `OPERATOR` as Vec<String>

* fix: std

Co-authored-by: Andrew Lamb <[email protected]>

Author: iskakaushik

src/ast/operator.rs

@@ -12,9 +12,13 @@
 
 use core::fmt;
 
+#[cfg(not(feature = "std"))]
+use alloc::{string::String, vec::Vec};
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
+use super::display_separated;
+
 /// Unary operators
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -86,41 +90,49 @@ pub enum BinaryOperator {
     PGRegexIMatch,
     PGRegexNotMatch,
     PGRegexNotIMatch,
+    /// PostgreSQL-specific custom operator.
+    ///
+    /// See [CREATE OPERATOR](https://www.postgresql.org/docs/current/sql-createoperator.html)
+    /// for more information.
+    PGCustomBinaryOperator(Vec<String>),
 }
 
 impl fmt::Display for BinaryOperator {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.write_str(match self {
-            BinaryOperator::Plus => "+",
-            BinaryOperator::Minus => "-",
-            BinaryOperator::Multiply => "*",
-            BinaryOperator::Divide => "/",
-            BinaryOperator::Modulo => "%",
-            BinaryOperator::StringConcat => "||",
-            BinaryOperator::Gt => ">",
-            BinaryOperator::Lt => "<",
-            BinaryOperator::GtEq => ">=",
-            BinaryOperator::LtEq => "<=",
-            BinaryOperator::Spaceship => "<=>",
-            BinaryOperator::Eq => "=",
-            BinaryOperator::NotEq => "<>",
-            BinaryOperator::And => "AND",
-            BinaryOperator::Or => "OR",
-            BinaryOperator::Xor => "XOR",
-            BinaryOperator::Like => "LIKE",
-            BinaryOperator::NotLike => "NOT LIKE",
-            BinaryOperator::ILike => "ILIKE",
-            BinaryOperator::NotILike => "NOT ILIKE",
-            BinaryOperator::BitwiseOr => "|",
-            BinaryOperator::BitwiseAnd => "&",
-            BinaryOperator::BitwiseXor => "^",
-            BinaryOperator::PGBitwiseXor => "#",
-            BinaryOperator::PGBitwiseShiftLeft => "<<",
-            BinaryOperator::PGBitwiseShiftRight => ">>",
-            BinaryOperator::PGRegexMatch => "~",
-            BinaryOperator::PGRegexIMatch => "~*",
-            BinaryOperator::PGRegexNotMatch => "!~",
-            BinaryOperator::PGRegexNotIMatch => "!~*",
-        })
+        match self {
+            BinaryOperator::Plus => f.write_str("+"),
+            BinaryOperator::Minus => f.write_str("-"),
+            BinaryOperator::Multiply => f.write_str("*"),
+            BinaryOperator::Divide => f.write_str("/"),
+            BinaryOperator::Modulo => f.write_str("%"),
+            BinaryOperator::StringConcat => f.write_str("||"),
+            BinaryOperator::Gt => f.write_str(">"),
+            BinaryOperator::Lt => f.write_str("<"),
+            BinaryOperator::GtEq => f.write_str(">="),
+            BinaryOperator::LtEq => f.write_str("<="),
+            BinaryOperator::Spaceship => f.write_str("<=>"),
+            BinaryOperator::Eq => f.write_str("="),
+            BinaryOperator::NotEq => f.write_str("<>"),
+            BinaryOperator::And => f.write_str("AND"),
+            BinaryOperator::Or => f.write_str("OR"),
+            BinaryOperator::Xor => f.write_str("XOR"),
+            BinaryOperator::Like => f.write_str("LIKE"),
+            BinaryOperator::NotLike => f.write_str("NOT LIKE"),
+            BinaryOperator::ILike => f.write_str("ILIKE"),
+            BinaryOperator::NotILike => f.write_str("NOT ILIKE"),
+            BinaryOperator::BitwiseOr => f.write_str("|"),
+            BinaryOperator::BitwiseAnd => f.write_str("&"),
+            BinaryOperator::BitwiseXor => f.write_str("^"),
+            BinaryOperator::PGBitwiseXor => f.write_str("#"),
+            BinaryOperator::PGBitwiseShiftLeft => f.write_str("<<"),
+            BinaryOperator::PGBitwiseShiftRight => f.write_str(">>"),
+            BinaryOperator::PGRegexMatch => f.write_str("~"),
+            BinaryOperator::PGRegexIMatch => f.write_str("~*"),
+            BinaryOperator::PGRegexNotMatch => f.write_str("!~"),
+            BinaryOperator::PGRegexNotIMatch => f.write_str("!~*"),
+            BinaryOperator::PGCustomBinaryOperator(idents) => {
+                write!(f, "OPERATOR({})", display_separated(idents, "."))
+            }
+        }
     }
 }

src/keywords.rs

@@ -358,6 +358,7 @@ define_keywords!(
     ON,
     ONLY,
     OPEN,
+    OPERATOR,
     OPTION,
     OR,
     ORC,

src/parser.rs

@@ -1173,6 +1173,22 @@ impl<'a> Parser<'a> {
                     }
                 }
                 Keyword::XOR => Some(BinaryOperator::Xor),
+                Keyword::OPERATOR if dialect_of!(self is PostgreSqlDialect | GenericDialect) => {
+                    self.expect_token(&Token::LParen)?;
+                    // there are special rules for operator names in
+                    // postgres so we can not use 'parse_object'
+                    // or similar.
+                    // See https://www.postgresql.org/docs/current/sql-createoperator.html
+                    let mut idents = vec![];
+                    loop {
+                        idents.push(self.next_token().to_string());
+                        if !self.consume_token(&Token::Period) {
+                            break;
+                        }
+                    }
+                    self.expect_token(&Token::RParen)?;
+                    Some(BinaryOperator::PGCustomBinaryOperator(idents))
+                }
                 _ => None,
             },
             _ => None,
@@ -1437,6 +1453,7 @@ impl<'a> Parser<'a> {
             Token::Word(w) if w.keyword == Keyword::BETWEEN => Ok(Self::BETWEEN_PREC),
             Token::Word(w) if w.keyword == Keyword::LIKE => Ok(Self::BETWEEN_PREC),
             Token::Word(w) if w.keyword == Keyword::ILIKE => Ok(Self::BETWEEN_PREC),
+            Token::Word(w) if w.keyword == Keyword::OPERATOR => Ok(Self::BETWEEN_PREC),
             Token::Eq
             | Token::Lt
             | Token::LtEq

tests/sqlparser_postgres.rs

@@ -1565,3 +1565,55 @@ fn parse_fetch() {
     pg_and_generic()
         .verified_stmt("FETCH BACKWARD ALL IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
 }
+
+#[test]
+fn parse_custom_operator() {
+    // operator with a database and schema
+    let sql = r#"SELECT * FROM events WHERE relname OPERATOR(database.pg_catalog.~) '^(table)$'"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        select.selection,
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident {
+                value: "relname".into(),
+                quote_style: None,
+            })),
+            op: BinaryOperator::PGCustomBinaryOperator(vec![
+                "database".into(),
+                "pg_catalog".into(),
+                "~".into()
+            ]),
+            right: Box::new(Expr::Value(Value::SingleQuotedString("^(table)$".into())))
+        })
+    );
+
+    // operator with a schema
+    let sql = r#"SELECT * FROM events WHERE relname OPERATOR(pg_catalog.~) '^(table)$'"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        select.selection,
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident {
+                value: "relname".into(),
+                quote_style: None,
+            })),
+            op: BinaryOperator::PGCustomBinaryOperator(vec!["pg_catalog".into(), "~".into()]),
+            right: Box::new(Expr::Value(Value::SingleQuotedString("^(table)$".into())))
+        })
+    );
+
+    // custom operator without a schema
+    let sql = r#"SELECT * FROM events WHERE relname OPERATOR(~) '^(table)$'"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        select.selection,
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident {
+                value: "relname".into(),
+                quote_style: None,
+            })),
+            op: BinaryOperator::PGCustomBinaryOperator(vec!["~".into()]),
+            right: Box::new(Expr::Value(Value::SingleQuotedString("^(table)$".into())))
+        })
+    );
+}