feat: add arg placeholder (apache#420)

Co-authored-by: gamife <[email protected]>

Author: 2 people

src/ast/value.rs

@@ -64,6 +64,8 @@ pub enum Value {
     },
     /// `NULL` value
     Null,
+    /// `?` or `$` Prepared statement arg placeholder
+    Placeholder(String),
 }
 
 impl fmt::Display for Value {
@@ -117,6 +119,7 @@ impl fmt::Display for Value {
                 Ok(())
             }
             Value::Null => write!(f, "NULL"),
+            Value::Placeholder(v) => write!(f, "{}", v),
         }
     }
 }

src/parser.rs

@@ -454,6 +454,10 @@ impl<'a> Parser<'a> {
                 self.expect_token(&Token::RParen)?;
                 Ok(expr)
             }
+            Token::Placeholder(_) => {
+                self.prev_token();
+                Ok(Expr::Value(self.parse_value()?))
+            }
             unexpected => self.expected("an expression:", unexpected),
         }?;
 
@@ -1991,6 +1995,7 @@ impl<'a> Parser<'a> {
             Token::SingleQuotedString(ref s) => Ok(Value::SingleQuotedString(s.to_string())),
             Token::NationalStringLiteral(ref s) => Ok(Value::NationalStringLiteral(s.to_string())),
             Token::HexStringLiteral(ref s) => Ok(Value::HexStringLiteral(s.to_string())),
+            Token::Placeholder(ref s) => Ok(Value::Placeholder(s.to_string())),
             unexpected => self.expected("a value", unexpected),
         }
     }

src/tokenizer.rs

@@ -139,6 +139,8 @@ pub enum Token {
     PGSquareRoot,
     /// `||/` , a cube root math operator in PostgreSQL
     PGCubeRoot,
+    /// `?` or `$` , a prepared statement arg placeholder
+    Placeholder(String),
 }
 
 impl fmt::Display for Token {
@@ -194,6 +196,7 @@ impl fmt::Display for Token {
             Token::ShiftRight => f.write_str(">>"),
             Token::PGSquareRoot => f.write_str("|/"),
             Token::PGCubeRoot => f.write_str("||/"),
+            Token::Placeholder(ref s) => write!(f, "{}", s),
         }
     }
 }
@@ -337,6 +340,7 @@ impl<'a> Tokenizer<'a> {
                 Token::Word(w) if w.quote_style != None => self.col += w.value.len() as u64 + 2,
                 Token::Number(s, _) => self.col += s.len() as u64,
                 Token::SingleQuotedString(s) => self.col += s.len() as u64,
+                Token::Placeholder(s) => self.col += s.len() as u64,
                 _ => self.col += 1,
             }
 
@@ -596,6 +600,15 @@ impl<'a> Tokenizer<'a> {
                 }
                 '#' => self.consume_and_return(chars, Token::Sharp),
                 '@' => self.consume_and_return(chars, Token::AtSign),
+                '?' => self.consume_and_return(chars, Token::Placeholder(String::from("?"))),
+                '$' => {
+                    chars.next();
+                    let s = peeking_take_while(
+                        chars,
+                        |ch| matches!(ch, '0'..='9' | 'A'..='Z' | 'a'..='z'),
+                    );
+                    Ok(Some(Token::Placeholder(String::from("$") + &s)))
+                }
                 other => self.consume_and_return(chars, Token::Char(other)),
             },
             None => Ok(None),

tests/sqlparser_common.rs

@@ -20,11 +20,14 @@
 
 #[macro_use]
 mod test_utils;
+use sqlparser::test_utils::TestedDialects;
 use test_utils::{all_dialects, expr_from_projection, join, number, only, table, table_alias};
 
 use matches::assert_matches;
 use sqlparser::ast::*;
-use sqlparser::dialect::{GenericDialect, SQLiteDialect};
+use sqlparser::dialect::{
+    AnsiDialect, GenericDialect, MsSqlDialect, PostgreSqlDialect, SQLiteDialect, SnowflakeDialect,
+};
 use sqlparser::keywords::ALL_KEYWORDS;
 use sqlparser::parser::{Parser, ParserError};
 
@@ -3693,6 +3696,42 @@ fn parse_drop_index() {
     }
 }
 
+#[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 {}),
+            Box::new(PostgreSqlDialect {}),
+            Box::new(MsSqlDialect {}),
+            Box::new(AnsiDialect {}),
+            Box::new(SnowflakeDialect {}),
+            // Note: `$` is the starting word for the HiveDialect identifier
+            // Box::new(sqlparser::dialect::HiveDialect {}),
+        ],
+    };
+    let sql = "SELECT * FROM student WHERE id = $Id1";
+    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("$Id1".into())))
+        })
+    );
+}
+
 #[test]
 fn all_keywords_sorted() {
     // assert!(ALL_KEYWORDS.is_sorted())