Add support for postgres json operators ->, ->>, #>, and #>> (#458)

* add support for postgres json selection

Signed-off-by: password <rbalajis25@gmail.com>

* fix clippy

Signed-off-by: password <rbalajis25@gmail.com>

* add support for postgres `#>` and `#>>` json operator

* fix clippy

Signed-off-by: poonai <rbalajis25@gmail.com>

* resolve comments

Signed-off-by: password <rbalajis25@gmail.com>

src/ast/mod.rs

@@ -180,6 +180,39 @@ impl fmt::Display for Array {
     }
 }
 
+/// JsonOperator
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum JsonOperator {
+    /// -> keeps the value as json
+    Arrow,
+    /// ->> keeps the value as text or int.
+    LongArrow,
+    /// #> Extracts JSON sub-object at the specified path
+    HashArrow,
+    /// #>> Extracts JSON sub-object at the specified path as text
+    HashLongArrow,
+}
+
+impl fmt::Display for JsonOperator {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            JsonOperator::Arrow => {
+                write!(f, "->")
+            }
+            JsonOperator::LongArrow => {
+                write!(f, "->>")
+            }
+            JsonOperator::HashArrow => {
+                write!(f, "#>")
+            }
+            JsonOperator::HashLongArrow => {
+                write!(f, "#>>")
+            }
+        }
+    }
+}
+
 /// An SQL expression of any type.
 ///
 /// The parser does not distinguish between expressions of different types
@@ -192,6 +225,12 @@ pub enum Expr {
     Identifier(Ident),
     /// Multi-part identifier, e.g. `table_alias.column` or `schema.table.col`
     CompoundIdentifier(Vec<Ident>),
+    /// JSON access (postgres)  eg: data->'tags'
+    JsonAccess {
+        left: Box<Expr>,
+        operator: JsonOperator,
+        right: Box<Expr>,
+    },
     /// `IS NULL` operator
     IsNull(Box<Expr>),
     /// `IS NOT NULL` operator
@@ -507,6 +546,13 @@ impl fmt::Display for Expr {
             Expr::Array(set) => {
                 write!(f, "{}", set)
             }
+            Expr::JsonAccess {
+                left,
+                operator,
+                right,
+            } => {
+                write!(f, "{} {} {}", left, operator, right)
+            }
         }
     }
 }

src/parser.rs

@@ -1146,6 +1146,23 @@ impl<'a> Parser<'a> {
                 return self.parse_array_index(expr);
             }
             self.parse_map_access(expr)
+        } else if Token::Arrow == tok
+            || Token::LongArrow == tok
+            || Token::HashArrow == tok
+            || Token::HashLongArrow == tok
+        {
+            let operator = match tok {
+                Token::Arrow => JsonOperator::Arrow,
+                Token::LongArrow => JsonOperator::LongArrow,
+                Token::HashArrow => JsonOperator::HashArrow,
+                Token::HashLongArrow => JsonOperator::HashLongArrow,
+                _ => unreachable!(),
+            };
+            Ok(Expr::JsonAccess {
+                left: Box::new(expr),
+                operator,
+                right: Box::new(self.parse_expr()?),
+            })
         } else {
             // Can only happen if `get_next_precedence` got out of sync with this function
             parser_err!(format!("No infix parser for token {:?}", tok))
@@ -1291,7 +1308,11 @@ impl<'a> Parser<'a> {
             Token::Mul | Token::Div | Token::Mod | Token::StringConcat => Ok(40),
             Token::DoubleColon => Ok(50),
             Token::ExclamationMark => Ok(50),
-            Token::LBracket => Ok(50),
+            Token::LBracket
+            | Token::LongArrow
+            | Token::Arrow
+            | Token::HashArrow
+            | Token::HashLongArrow => Ok(50),
             _ => Ok(0),
         }
     }

src/tokenizer.rs

@@ -141,6 +141,14 @@ pub enum Token {
     PGCubeRoot,
     /// `?` or `$` , a prepared statement arg placeholder
     Placeholder(String),
+    /// ->, used as a operator to extract json field in PostgreSQL
+    Arrow,
+    /// ->>, used as a operator to extract json field as text in PostgreSQL
+    LongArrow,
+    /// #> Extracts JSON sub-object at the specified path
+    HashArrow,
+    /// #>> Extracts JSON sub-object at the specified path as text
+    HashLongArrow,
 }
 
 impl fmt::Display for Token {
@@ -197,6 +205,10 @@ impl fmt::Display for Token {
             Token::PGSquareRoot => f.write_str("|/"),
             Token::PGCubeRoot => f.write_str("||/"),
             Token::Placeholder(ref s) => write!(f, "{}", s),
+            Token::Arrow => write!(f, "->"),
+            Token::LongArrow => write!(f, "->>"),
+            Token::HashArrow => write!(f, "#>"),
+            Token::HashLongArrow => write!(f, "#>>"),
         }
     }
 }
@@ -483,6 +495,16 @@ impl<'a> Tokenizer<'a> {
                                 comment,
                             })))
                         }
+                        Some('>') => {
+                            chars.next();
+                            match chars.peek() {
+                                Some('>') => {
+                                    chars.next();
+                                    Ok(Some(Token::LongArrow))
+                                }
+                                _ => Ok(Some(Token::Arrow)),
+                            }
+                        }
                         // a regular '-' operator
                         _ => Ok(Some(Token::Minus)),
                     }
@@ -600,7 +622,22 @@ impl<'a> Tokenizer<'a> {
                         _ => Ok(Some(Token::Tilde)),
                     }
                 }
-                '#' => self.consume_and_return(chars, Token::Sharp),
+                '#' => {
+                    chars.next();
+                    match chars.peek() {
+                        Some('>') => {
+                            chars.next();
+                            match chars.peek() {
+                                Some('>') => {
+                                    chars.next();
+                                    Ok(Some(Token::HashLongArrow))
+                                }
+                                _ => Ok(Some(Token::HashArrow)),
+                            }
+                        }
+                        _ => Ok(Some(Token::Sharp)),
+                    }
+                }
                 '@' => self.consume_and_return(chars, Token::AtSign),
                 '?' => self.consume_and_return(chars, Token::Placeholder(String::from("?"))),
                 '$' => {

tests/sqlparser_postgres.rs

@@ -1233,6 +1233,74 @@ fn test_savepoint() {
     }
 }
 
+#[test]
+fn test_json() {
+    let sql = "SELECT params ->> 'name' FROM events";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("params"))),
+            operator: JsonOperator::LongArrow,
+            right: Box::new(Expr::Value(Value::SingleQuotedString("name".to_string()))),
+        }),
+        select.projection[0]
+    );
+
+    let sql = "SELECT params -> 'name' FROM events";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("params"))),
+            operator: JsonOperator::Arrow,
+            right: Box::new(Expr::Value(Value::SingleQuotedString("name".to_string()))),
+        }),
+        select.projection[0]
+    );
+
+    let sql = "SELECT info -> 'items' ->> 'product' FROM orders";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            operator: JsonOperator::Arrow,
+            right: Box::new(Expr::JsonAccess {
+                left: Box::new(Expr::Value(Value::SingleQuotedString("items".to_string()))),
+                operator: JsonOperator::LongArrow,
+                right: Box::new(Expr::Value(Value::SingleQuotedString(
+                    "product".to_string()
+                )))
+            }),
+        }),
+        select.projection[0]
+    );
+
+    let sql = "SELECT info #> '{a,b,c}' FROM orders";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            operator: JsonOperator::HashArrow,
+            right: Box::new(Expr::Value(Value::SingleQuotedString(
+                "{a,b,c}".to_string()
+            ))),
+        }),
+        select.projection[0]
+    );
+
+    let sql = "SELECT info #>> '{a,b,c}' FROM orders";
+    let select = pg().verified_only_select(sql);
+    assert_eq!(
+        SelectItem::UnnamedExpr(Expr::JsonAccess {
+            left: Box::new(Expr::Identifier(Ident::new("info"))),
+            operator: JsonOperator::HashLongArrow,
+            right: Box::new(Expr::Value(Value::SingleQuotedString(
+                "{a,b,c}".to_string()
+            ))),
+        }),
+        select.projection[0]
+    );
+}
+
 #[test]
 fn parse_comments() {
     match pg().verified_stmt("COMMENT ON COLUMN tab.name IS 'comment'") {