Support ?-based jsonb operators in Postgres (apache#1242)

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

Author: ReppCodes

src/ast/operator.rs

@@ -214,6 +214,27 @@ pub enum BinaryOperator {
     ///
     /// See <https://www.postgresql.org/docs/current/functions-json.html>.
     AtQuestion,
+    /// The `?` operator.
+    ///
+    /// On PostgreSQL, this operator is used to check whether a string exists as a top-level key
+    /// within the JSON value
+    ///
+    /// See <https://www.postgresql.org/docs/current/functions-json.html>.
+    Question,
+    /// The `?&` operator.
+    ///
+    /// On PostgreSQL, this operator is used to check whether all of the the indicated array
+    /// members exist as top-level keys.
+    ///
+    /// See <https://www.postgresql.org/docs/current/functions-json.html>.
+    QuestionAnd,
+    /// The `?|` operator.
+    ///
+    /// On PostgreSQL, this operator is used to check whether any of the the indicated array
+    /// members exist as top-level keys.
+    ///
+    /// See <https://www.postgresql.org/docs/current/functions-json.html>.
+    QuestionPipe,
     /// PostgreSQL-specific custom operator.
     ///
     /// See [CREATE OPERATOR](https://www.postgresql.org/docs/current/sql-createoperator.html)
@@ -269,6 +290,9 @@ impl fmt::Display for BinaryOperator {
             BinaryOperator::ArrowAt => f.write_str("<@"),
             BinaryOperator::HashMinus => f.write_str("#-"),
             BinaryOperator::AtQuestion => f.write_str("@?"),
+            BinaryOperator::Question => f.write_str("?"),
+            BinaryOperator::QuestionAnd => f.write_str("?&"),
+            BinaryOperator::QuestionPipe => f.write_str("?|"),
             BinaryOperator::PGCustomBinaryOperator(idents) => {
                 write!(f, "OPERATOR({})", display_separated(idents, "."))
             }

src/parser/mod.rs

@@ -2382,6 +2382,9 @@ impl<'a> Parser<'a> {
             Token::HashMinus => Some(BinaryOperator::HashMinus),
             Token::AtQuestion => Some(BinaryOperator::AtQuestion),
             Token::AtAt => Some(BinaryOperator::AtAt),
+            Token::Question => Some(BinaryOperator::Question),
+            Token::QuestionAnd => Some(BinaryOperator::QuestionAnd),
+            Token::QuestionPipe => Some(BinaryOperator::QuestionPipe),
 
             Token::Word(w) => match w.keyword {
                 Keyword::AND => Some(BinaryOperator::And),
@@ -2878,7 +2881,10 @@ impl<'a> Parser<'a> {
             | Token::ArrowAt
             | Token::HashMinus
             | Token::AtQuestion
-            | Token::AtAt => Ok(Self::PG_OTHER_PREC),
+            | Token::AtAt
+            | Token::Question
+            | Token::QuestionAnd
+            | Token::QuestionPipe => Ok(Self::PG_OTHER_PREC),
             _ => Ok(0),
         }
     }

src/tokenizer.rs

@@ -36,7 +36,8 @@ use sqlparser_derive::{Visit, VisitMut};
 
 use crate::ast::DollarQuotedString;
 use crate::dialect::{
-    BigQueryDialect, DuckDbDialect, GenericDialect, HiveDialect, SnowflakeDialect,
+    BigQueryDialect, DuckDbDialect, GenericDialect, HiveDialect, PostgreSqlDialect,
+    SnowflakeDialect,
 };
 use crate::dialect::{Dialect, MySqlDialect};
 use crate::keywords::{Keyword, ALL_KEYWORDS, ALL_KEYWORDS_INDEX};
@@ -199,6 +200,15 @@ pub enum Token {
     /// for the specified JSON value. Only the first item of the result is taken into
     /// account. If the result is not Boolean, then NULL is returned.
     AtAt,
+    /// jsonb ? text -> boolean: Checks whether the string exists as a top-level key within the
+    /// jsonb object
+    Question,
+    /// jsonb ?& text[] -> boolean: Check whether all members of the text array exist as top-level
+    /// keys within the jsonb object
+    QuestionAnd,
+    /// jsonb ?| text[] -> boolean: Check whether any member of the text array exists as top-level
+    /// keys within the jsonb object
+    QuestionPipe,
 }
 
 impl fmt::Display for Token {
@@ -278,6 +288,9 @@ impl fmt::Display for Token {
             Token::HashMinus => write!(f, "#-"),
             Token::AtQuestion => write!(f, "@?"),
             Token::AtAt => write!(f, "@@"),
+            Token::Question => write!(f, "?"),
+            Token::QuestionAnd => write!(f, "?&"),
+            Token::QuestionPipe => write!(f, "?|"),
         }
     }
 }
@@ -1059,6 +1072,15 @@ impl<'a> Tokenizer<'a> {
                         _ => Ok(Some(Token::AtSign)),
                     }
                 }
+                // Postgres uses ? for jsonb operators, not prepared statements
+                '?' if dialect_of!(self is PostgreSqlDialect) => {
+                    chars.next();
+                    match chars.peek() {
+                        Some('|') => self.consume_and_return(chars, Token::QuestionPipe),
+                        Some('&') => self.consume_and_return(chars, Token::QuestionAnd),
+                        _ => self.consume_and_return(chars, Token::Question),
+                    }
+                }
                 '?' => {
                     chars.next();
                     let s = peeking_take_while(chars, |ch| ch.is_numeric());

tests/sqlparser_common.rs

@@ -7802,17 +7802,6 @@ fn test_lock_nonblock() {
 
 #[test]
 fn test_placeholder() {
-    let sql = "SELECT * FROM student WHERE id = ?";
-    let ast = verified_only_select(sql);
-    assert_eq!(
-        ast.selection,
-        Some(Expr::BinaryOp {
-            left: Box::new(Expr::Identifier(Ident::new("id"))),
-            op: BinaryOperator::Eq,
-            right: Box::new(Expr::Value(Value::Placeholder("?".into()))),
-        })
-    );
-
     let dialects = TestedDialects {
         dialects: vec![
             Box::new(GenericDialect {}),
@@ -7852,6 +7841,32 @@ fn test_placeholder() {
         }),
     );
 
+    let dialects = TestedDialects {
+        dialects: vec![
+            Box::new(GenericDialect {}),
+            Box::new(DuckDbDialect {}),
+            // Note: `?` is for jsonb operators in PostgreSqlDialect
+            // Box::new(PostgreSqlDialect {}),
+            Box::new(MsSqlDialect {}),
+            Box::new(AnsiDialect {}),
+            Box::new(BigQueryDialect {}),
+            Box::new(SnowflakeDialect {}),
+            // Note: `$` is the starting word for the HiveDialect identifier
+            // Box::new(sqlparser::dialect::HiveDialect {}),
+        ],
+        options: None,
+    };
+    let sql = "SELECT * FROM student WHERE id = ?";
+    let ast = dialects.verified_only_select(sql);
+    assert_eq!(
+        ast.selection,
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("id"))),
+            op: BinaryOperator::Eq,
+            right: Box::new(Expr::Value(Value::Placeholder("?".into()))),
+        })
+    );
+
     let sql = "SELECT $fromage_français, :x, ?123";
     let ast = dialects.verified_only_select(sql);
     assert_eq!(

tests/sqlparser_postgres.rs

@@ -2401,6 +2401,51 @@ fn test_json() {
         },
         select.selection.unwrap(),
     );
+
+    let sql = r#"SELECT info FROM orders WHERE info ? 'b'"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            op: BinaryOperator::Question,
+            right: Box::new(Expr::Value(Value::SingleQuotedString("b".to_string()))),
+        },
+        select.selection.unwrap(),
+    );
+
+    let sql = r#"SELECT info FROM orders WHERE info ?& ARRAY['b', 'c']"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            op: BinaryOperator::QuestionAnd,
+            right: Box::new(Expr::Array(Array {
+                elem: vec![
+                    Expr::Value(Value::SingleQuotedString("b".to_string())),
+                    Expr::Value(Value::SingleQuotedString("c".to_string()))
+                ],
+                named: true
+            }))
+        },
+        select.selection.unwrap(),
+    );
+
+    let sql = r#"SELECT info FROM orders WHERE info ?| ARRAY['b', 'c']"#;
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            op: BinaryOperator::QuestionPipe,
+            right: Box::new(Expr::Array(Array {
+                elem: vec![
+                    Expr::Value(Value::SingleQuotedString("b".to_string())),
+                    Expr::Value(Value::SingleQuotedString("c".to_string()))
+                ],
+                named: true
+            }))
+        },
+        select.selection.unwrap(),
+    );
 }
 
 #[test]