Allow alphanumeric characters in SQLLite style parameters. (#32)

Author: BlurrechDev

src/lib.rs

@@ -16,7 +16,9 @@ mod tokenizer;
 /// Formats whitespace in a SQL string to make it easier to read.
 /// Optionally replaces parameter placeholders with `params`.
 pub fn format(query: &str, params: &QueryParams, options: FormatOptions) -> String {
-    let tokens = tokenizer::tokenize(query);
+    let named_placeholders = matches!(params, QueryParams::Named(_));
+
+    let tokens = tokenizer::tokenize(query, named_placeholders);
     formatter::format(&tokens, params, options)
 }
 
@@ -1067,21 +1069,40 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, options), expected);
     }
 
+    #[test]
+    fn it_recognizes_dollar_sign_alphanumeric_placeholders() {
+        let input = "SELECT $hash, $foo, $bar;";
+        let options = FormatOptions::default();
+        let expected = indoc!(
+            "
+            SELECT
+              $hash,
+              $foo,
+              $bar;"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, options), expected);
+    }
+
     #[test]
     fn it_recognizes_dollar_sign_numbered_placeholders_with_param_values() {
-        let input = "SELECT $2, $3, $1;";
+        let input = "SELECT $2, $3, $1, $named, $4, $alias;";
         let params = vec![
             "first".to_string(),
             "second".to_string(),
             "third".to_string(),
+            "4th".to_string(),
         ];
         let options = FormatOptions::default();
         let expected = indoc!(
             "
             SELECT
               second,
               third,
-              first;"
+              first,
+              $named,
+              4th,
+              $alias;"
         );
 
         assert_eq!(
@@ -1090,6 +1111,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn it_recognizes_dollar_sign_alphanumeric_placeholders_with_param_values() {
+        let input = "SELECT $hash, $salt, $1, $2;";
+        let params = vec![
+            ("hash".to_string(), "hash value".to_string()),
+            ("salt".to_string(), "salt value".to_string()),
+            ("1".to_string(), "number 1".to_string()),
+            ("2".to_string(), "number 2".to_string()),
+        ];
+        let options = FormatOptions::default();
+        let expected = indoc!(
+            "
+            SELECT
+              hash value,
+              salt value,
+              number 1,
+              number 2;"
+        );
+
+        assert_eq!(
+            format(input, &QueryParams::Named(params), options),
+            expected
+        );
+    }
+
     #[test]
     fn it_formats_query_with_go_batch_separator() {
         let input = "SELECT 1 GO SELECT 2";

src/params.rs

@@ -12,6 +12,8 @@ impl<'a> Params<'a> {
     }
 
     pub fn get(&mut self, token: &'a Token<'a>) -> &'a str {
+        let named_placeholder_token = token.key.as_ref().map_or(false, |key| key.named() != "");
+
         match self.params {
             QueryParams::Named(params) => token
                 .key
@@ -23,7 +25,7 @@ impl<'a> Params<'a> {
                         .map(|param| param.1.as_str())
                 })
                 .unwrap_or(token.value),
-            QueryParams::Indexed(params) => {
+            QueryParams::Indexed(params) if !named_placeholder_token => {
                 if let Some(key) = token.key.as_ref().and_then(|key| key.indexed()) {
                     params
                         .get(key)
@@ -38,7 +40,7 @@ impl<'a> Params<'a> {
                     value
                 }
             }
-            QueryParams::None => token.value,
+            _ => token.value,
         }
     }
 }

src/tokenizer.rs

@@ -10,15 +10,18 @@ use nom::{AsChar, Err, IResult};
 use std::borrow::Cow;
 use unicode_categories::UnicodeCategories;
 
-pub(crate) fn tokenize(mut input: &str) -> Vec<Token<'_>> {
+pub(crate) fn tokenize(mut input: &str, named_placeholders: bool) -> Vec<Token<'_>> {
     let mut tokens: Vec<Token> = Vec::new();
 
     let mut last_reserved_token = None;
 
     // Keep processing the string until it is empty
-    while let Ok(result) =
-        get_next_token(input, tokens.last().cloned(), last_reserved_token.clone())
-    {
+    while let Ok(result) = get_next_token(
+        input,
+        tokens.last().cloned(),
+        last_reserved_token.clone(),
+        named_placeholders,
+    ) {
         if result.1.kind == TokenKind::Reserved {
             last_reserved_token = Some(result.1.clone());
         }
@@ -83,15 +86,16 @@ fn get_next_token<'a>(
     input: &'a str,
     previous_token: Option<Token<'a>>,
     last_reserved_token: Option<Token<'a>>,
+    named_placeholders: bool,
 ) -> IResult<&'a str, Token<'a>> {
     get_whitespace_token(input)
         .or_else(|_| get_comment_token(input))
         .or_else(|_| get_string_token(input))
         .or_else(|_| get_open_paren_token(input))
         .or_else(|_| get_close_paren_token(input))
-        .or_else(|_| get_placeholder_token(input))
         .or_else(|_| get_number_token(input))
         .or_else(|_| get_reserved_word_token(input, previous_token, last_reserved_token))
+        .or_else(|_| get_placeholder_token(input, named_placeholders))
         .or_else(|_| get_word_token(input))
         .or_else(|_| get_operator_token(input))
 }
@@ -288,12 +292,23 @@ fn get_close_paren_token(input: &str) -> IResult<&str, Token<'_>> {
     })
 }
 
-fn get_placeholder_token(input: &str) -> IResult<&str, Token<'_>> {
-    alt((
-        get_ident_named_placeholder_token,
-        get_string_named_placeholder_token,
-        get_indexed_placeholder_token,
-    ))(input)
+fn get_placeholder_token(input: &str, named_placeholders: bool) -> IResult<&str, Token<'_>> {
+    // The precedence changes based on 'named_placeholders' but not the exhaustiveness.
+    // This is to ensure the formatting is the same even if parameters aren't used.
+
+    if named_placeholders {
+        alt((
+            get_ident_named_placeholder_token,
+            get_string_named_placeholder_token,
+            get_indexed_placeholder_token,
+        ))(input)
+    } else {
+        alt((
+            get_indexed_placeholder_token,
+            get_ident_named_placeholder_token,
+            get_string_named_placeholder_token,
+        ))(input)
+    }
 }
 
 fn get_indexed_placeholder_token(input: &str) -> IResult<&str, Token<'_>> {
@@ -327,7 +342,7 @@ fn get_indexed_placeholder_token(input: &str) -> IResult<&str, Token<'_>> {
 
 fn get_ident_named_placeholder_token(input: &str) -> IResult<&str, Token<'_>> {
     recognize(tuple((
-        alt((char('@'), char(':'))),
+        alt((char('@'), char(':'), char('$'))),
         take_while1(|item: char| {
             item.is_alphanumeric() || item == '.' || item == '_' || item == '$'
         }),