mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-22 15:04:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

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

Browse source 
Co-authored-by: Andrew Repp <arepp@cloudflare.com>
This commit is contained in:
Andrew Repp 2024-05-01 06:50:45 -04:00 • committed by GitHub
parent bafaa914b0
commit 4aa37a46a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 125 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

24
src/ast/operator.rs
View file

@ -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, "."))
}

8
src/parser/mod.rs
View file

@ -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),
}
}

24
src/tokenizer.rs
View file

@ -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());

37
tests/sqlparser_common.rs
View file

@ -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!(

45
tests/sqlparser_postgres.rs
View file

@ -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]