Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(parser): support to parse param in parser #8113

Merged
merged 3 commits into from
Feb 26, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/sqlparser/src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,8 @@ pub enum Expr {
Nested(Box<Expr>),
/// A literal value, such as string, number, date or NULL
Value(Value),
/// Parameter Symbol e.g. `$1`, `$1::int`
Parameter { index: u64 },
/// A constant of form `<data_type> 'value'`.
/// This can represent ANSI SQL `DATE`, `TIME`, and `TIMESTAMP` literals (such as `DATE
/// '2020-01-01'`), as well as constants of other types (a non-standard PostgreSQL extension).
Expand Down Expand Up @@ -465,6 +467,7 @@ impl fmt::Display for Expr {
Expr::Collate { expr, collation } => write!(f, "{} COLLATE {}", expr, collation),
Expr::Nested(ast) => write!(f, "({})", ast),
Expr::Value(v) => write!(f, "{}", v),
Expr::Parameter { index } => write!(f, "${}", index),
Expr::TypedString { data_type, value } => {
write!(f, "{}", data_type)?;
write!(f, " '{}'", &value::escape_single_quote_string(value))
Expand Down
10 changes: 9 additions & 1 deletion src/sqlparser/src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -569,7 +569,7 @@ impl Parser {
self.prev_token();
Ok(Expr::Value(self.parse_value()?))
}

Token::Parameter(number) => self.parse_param(number),
Token::LParen => {
let expr =
if self.parse_keyword(Keyword::SELECT) || self.parse_keyword(Keyword::WITH) {
Expand Down Expand Up @@ -603,6 +603,14 @@ impl Parser {
}
}

fn parse_param(&mut self, param: String) -> Result<Expr, ParserError> {
Ok(Expr::Parameter {
index: param.parse().map_err(|_| {
ParserError::ParserError(format!("Parameter symbol has a invalid index {}.", param))
})?,
})
}

/// Parses a field selection expression. See also [`Expr::FieldIdentifier`].
pub fn parse_struct_selection(&mut self, expr: Expr) -> Result<Expr, ParserError> {
let mut nested_expr = expr;
Expand Down
22 changes: 21 additions & 1 deletion src/sqlparser/src/tokenizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ pub enum Token {
NationalStringLiteral(String),
/// Hexadecimal string literal: i.e.: X'deadbeef'
HexStringLiteral(String),
/// Parameter symbols: i.e: $1, $2
Parameter(String),
/// Comma
Comma,
/// Whitespace (space, tab, etc)
Expand Down Expand Up @@ -161,6 +163,7 @@ impl fmt::Display for Token {
Token::NationalStringLiteral(ref s) => write!(f, "N'{}'", s),
Token::HexStringLiteral(ref s) => write!(f, "X'{}'", s),
Token::CstyleEscapesString(ref s) => write!(f, "E'{}'", s),
Token::Parameter(ref s) => write!(f, "${}", s),
Token::Comma => f.write_str(","),
Token::Whitespace(ws) => write!(f, "{}", ws),
Token::DoubleEq => f.write_str("=="),
Expand Down Expand Up @@ -359,7 +362,6 @@ impl<'a> Tokenizer<'a> {

/// Get the next token or return None
fn next_token(&self, chars: &mut Peekable<Chars<'_>>) -> Result<Option<Token>, TokenizerError> {
// println!("next_token: {:?}", chars.peek());
match chars.peek() {
Some(&ch) => match ch {
' ' => self.consume_and_return(chars, Token::Whitespace(Whitespace::Space)),
Expand Down Expand Up @@ -611,6 +613,13 @@ impl<'a> Tokenizer<'a> {
_ => Ok(Some(Token::Colon)),
}
}
'$' => {
if let Some(parameter) = self.tokenize_parameter(chars) {
Ok(Some(parameter))
} else {
Ok(Some(Token::Char('$')))
}
}
Comment on lines +616 to +622
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI in PostgreSQL $ can also start a dollar-quoted string constant (upstream sqlparser-rs/sqlparser-rs#772).

But extending the support later would not break common use cases.

';' => self.consume_and_return(chars, Token::SemiColon),
'\\' => self.consume_and_return(chars, Token::Backslash),
'[' => self.consume_and_return(chars, Token::LBracket),
Expand Down Expand Up @@ -790,6 +799,17 @@ impl<'a> Tokenizer<'a> {
chars.next();
Ok(Some(t))
}

fn tokenize_parameter(&self, chars: &mut Peekable<Chars<'_>>) -> Option<Token> {
chars.next(); // consume '$'

let s = peeking_take_while(chars, |ch| ch.is_ascii_digit());
if s.is_empty() {
None
} else {
Some(Token::Parameter(s))
}
}
}

/// Read from `chars` until `predicate` returns `false` or EOF is hit.
Expand Down
50 changes: 50 additions & 0 deletions src/sqlparser/tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1049,3 +1049,53 @@ fn parse_array() {
)),
);
}

#[test]
fn parse_param_symbol() {
let select = verified_only_select("SELECT $1");
assert_eq!(
SelectItem::UnnamedExpr(Expr::Parameter { index: 1 }),
select.projection[0]
);

let select = verified_only_select("SELECT *, $2 FROM t WHERE a = $1");
assert_eq!(
Expr::BinaryOp {
left: Box::new(Expr::Identifier(Ident::new_unchecked("a"))),
op: BinaryOperator::Eq,
right: Box::new(Expr::Parameter { index: 1 })
},
select.selection.unwrap()
);
assert_eq!(
SelectItem::UnnamedExpr(Expr::Parameter { index: 2 }),
select.projection[1]
);

let select = verified_only_select("SELECT CAST($4096 AS INT)");
assert_eq!(
SelectItem::UnnamedExpr(Expr::Cast {
expr: Box::new(Expr::Parameter { index: 4096 }),
data_type: DataType::Int
}),
select.projection[0]
);

let select = verified_only_select("SELECT * FROM t WHERE a = CAST($1024 AS BIGINT)");
assert_eq!(
Expr::BinaryOp {
left: Box::new(Expr::Identifier(Ident::new_unchecked("a"))),
op: BinaryOperator::Eq,
right: Box::new(Expr::Cast {
expr: Box::new(Expr::Parameter { index: 1024 }),
data_type: DataType::BigInt
})
},
select.selection.unwrap()
);

let query = verified_query("VALUES ($1)");
if let SetExpr::Values(values) = query.body {
assert_eq!(values.0[0][0], Expr::Parameter { index: 1 });
}
}