Support json operators @> <@ #- @? and @@

postgres supports a bunch more json operators. See https://www.postgresql.org/docs/15/functions-json.html Skipping operators starting with a question mark for now, since those are hard to distinguish from placeholders without more context.
sqlparser-rs · Dec 14, 2022 · 6d6eb4b · 6d6eb4b
1 parent fb02344
commit 6d6eb4b
Show file tree

Hide file tree

Showing 4 changed files with 132 additions and 2 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -190,6 +190,20 @@ pub enum JsonOperator {
     HashLongArrow,
     /// : Colon is used by Snowflake (Which is similar to LongArrow)
     Colon,
+    /// jsonb @> jsonb -> boolean: Test whether left json contains the right json
+    AtArrow,
+    /// jsonb <@ jsonb -> boolean: Test whether right json contains the left json
+    ArrowAt,
+    /// jsonb #- text[] -> jsonb: Deletes the field or array element at the specified
+    /// path, where path elements can be either field keys or array indexes.
+    HashMinus,
+    /// jsonb @? jsonpath -> boolean: Does JSON path return any item for the specified
+    /// JSON value?
+    AtQuestion,
+    /// jsonb @@ jsonpath → boolean: Returns the result of a JSON path predicate check
+    /// 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,
 }
 
 impl fmt::Display for JsonOperator {
@@ -210,6 +224,13 @@ impl fmt::Display for JsonOperator {
             JsonOperator::Colon => {
                 write!(f, ":")
             }
+            JsonOperator::AtArrow => {
+                write!(f, "@>")
+            }
+            JsonOperator::ArrowAt => write!(f, "<@"),
+            JsonOperator::HashMinus => write!(f, "#-"),
+            JsonOperator::AtQuestion => write!(f, "@?"),
+            JsonOperator::AtAt => write!(f, "@@"),
         }
     }
 }

diff --git a/src/parser.rs b/src/parser.rs
@@ -1602,12 +1602,22 @@ impl<'a> Parser<'a> {
             || Token::LongArrow == tok
             || Token::HashArrow == tok
             || Token::HashLongArrow == tok
+            || Token::AtArrow == tok
+            || Token::ArrowAt == tok
+            || Token::HashMinus == tok
+            || Token::AtQuestion == tok
+            || Token::AtAt == tok
         {
             let operator = match tok.token {
                 Token::Arrow => JsonOperator::Arrow,
                 Token::LongArrow => JsonOperator::LongArrow,
                 Token::HashArrow => JsonOperator::HashArrow,
                 Token::HashLongArrow => JsonOperator::HashLongArrow,
+                Token::AtArrow => JsonOperator::AtArrow,
+                Token::ArrowAt => JsonOperator::ArrowAt,
+                Token::HashMinus => JsonOperator::HashMinus,
+                Token::AtQuestion => JsonOperator::AtQuestion,
+                Token::AtAt => JsonOperator::AtAt,
                 _ => unreachable!(),
             };
             Ok(Expr::JsonAccess {
@@ -1805,7 +1815,12 @@ impl<'a> Parser<'a> {
             | Token::LongArrow
             | Token::Arrow
             | Token::HashArrow
-            | Token::HashLongArrow => Ok(50),
+            | Token::HashLongArrow
+            | Token::AtArrow
+            | Token::ArrowAt
+            | Token::HashMinus
+            | Token::AtQuestion
+            | Token::AtAt => Ok(50),
             _ => Ok(0),
         }
     }

diff --git a/src/tokenizer.rs b/src/tokenizer.rs
@@ -155,6 +155,20 @@ pub enum Token {
     HashArrow,
     /// #>> Extracts JSON sub-object at the specified path as text
     HashLongArrow,
+    /// jsonb @> jsonb -> boolean: Test whether left json contains the right json
+    AtArrow,
+    /// jsonb <@ jsonb -> boolean: Test whether right json contains the left json
+    ArrowAt,
+    /// jsonb #- text[] -> jsonb: Deletes the field or array element at the specified
+    /// path, where path elements can be either field keys or array indexes.
+    HashMinus,
+    /// jsonb @? jsonpath -> boolean: Does JSON path return any item for the specified
+    /// JSON value?
+    AtQuestion,
+    /// jsonb @@ jsonpath → boolean: Returns the result of a JSON path predicate check
+    /// 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,
 }
 
 impl fmt::Display for Token {
@@ -217,7 +231,12 @@ impl fmt::Display for Token {
             Token::LongArrow => write!(f, "->>"),
             Token::HashArrow => write!(f, "#>"),
             Token::HashLongArrow => write!(f, "#>>"),
+            Token::AtArrow => write!(f, "@>"),
             Token::DoubleDollarQuoting => write!(f, "$$"),
+            Token::ArrowAt => write!(f, "<@"),
+            Token::HashMinus => write!(f, "#-"),
+            Token::AtQuestion => write!(f, "@?"),
+            Token::AtAt => write!(f, "@@"),
         }
     }
 }
@@ -708,6 +727,7 @@ impl<'a> Tokenizer<'a> {
                         }
                         Some('>') => self.consume_and_return(chars, Token::Neq),
                         Some('<') => self.consume_and_return(chars, Token::ShiftLeft),
+                        Some('@') => self.consume_and_return(chars, Token::ArrowAt),
                         _ => Ok(Some(Token::Lt)),
                     }
                 }
@@ -752,6 +772,7 @@ impl<'a> Tokenizer<'a> {
                 '#' => {
                     chars.next();
                     match chars.peek() {
+                        Some('-') => self.consume_and_return(chars, Token::HashMinus),
                         Some('>') => {
                             chars.next();
                             match chars.peek() {
@@ -765,7 +786,15 @@ impl<'a> Tokenizer<'a> {
                         _ => Ok(Some(Token::Sharp)),
                     }
                 }
-                '@' => self.consume_and_return(chars, Token::AtSign),
+                '@' => {
+                    chars.next();
+                    match chars.peek() {
+                        Some('>') => self.consume_and_return(chars, Token::AtArrow),
+                        Some('?') => self.consume_and_return(chars, Token::AtQuestion),
+                        Some('@') => self.consume_and_return(chars, Token::AtAt),
+                        _ => Ok(Some(Token::AtSign)),
+                    }
+                }
                 '?' => {
                     chars.next();
                     let s = peeking_take_while(chars, |ch| ch.is_numeric());

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -1620,6 +1620,71 @@ fn test_json() {
         }),
         select.projection[0]
     );
+
+    let sql = "SELECT info FROM orders WHERE info @> '{\"a\": 1}'";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            operator: JsonOperator::AtArrow,
+            right: Box::new(Expr::Value(Value::SingleQuotedString(
+                "{\"a\": 1}".to_string()
+            ))),
+        },
+        select.selection.unwrap(),
+    );
+
+    let sql = "SELECT info FROM orders WHERE '{\"a\": 1}' <@ info";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::JsonAccess {
+            left: Box::new(Expr::Value(Value::SingleQuotedString(
+                "{\"a\": 1}".to_string()
+            ))),
+            operator: JsonOperator::ArrowAt,
+            right: Box::new(Expr::Identifier(Ident::new("info"))),
+        },
+        select.selection.unwrap(),
+    );
+
+    let sql = "SELECT info #- ARRAY['a', 'b'] FROM orders";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::from("info"))),
+            operator: JsonOperator::HashMinus,
+            right: Box::new(Expr::Array(Array {
+                elem: vec![
+                    Expr::Value(Value::SingleQuotedString("a".to_string())),
+                    Expr::Value(Value::SingleQuotedString("b".to_string())),
+                ],
+                named: true,
+            })),
+        }),
+        select.projection[0],
+    );
+
+    let sql = "SELECT info FROM orders WHERE info @? '$.a'";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::from("info"))),
+            operator: JsonOperator::AtQuestion,
+            right: Box::new(Expr::Value(Value::SingleQuotedString("$.a".to_string())),),
+        },
+        select.selection.unwrap(),
+    );
+
+    let sql = "SELECT info FROM orders WHERE info @@ '$.a'";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::from("info"))),
+            operator: JsonOperator::AtAt,
+            right: Box::new(Expr::Value(Value::SingleQuotedString("$.a".to_string())),),
+        },
+        select.selection.unwrap(),
+    );
 }
 
 #[test]