fix: Possible to provide an escape hatch for expressions (#51)

Cargo.toml

@@ -15,6 +15,8 @@ categories = ["development-tools"]
 [dependencies]
 nom = "7.0.0"
 unicode_categories = "0.1.1"
+once_cell = "1"
+regex = "=1.6"
 
 [dev-dependencies]
 criterion = "0.4"

src/formatter.rs

@@ -1,3 +1,4 @@
+use regex::Regex;
 use std::borrow::Cow;
 
 use crate::indentation::Indentation;
@@ -6,12 +7,40 @@ use crate::params::Params;
 use crate::tokenizer::{Token, TokenKind};
 use crate::{FormatOptions, QueryParams};
 
+use once_cell::sync::Lazy;
+
+static RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i)^(--|/\*)\s*fmt\s*:\s*(off|on)").unwrap());
+
+pub(crate) fn check_fmt_off(s: &str) -> Option<bool> {
+    RE.captures(s)?
+        .get(2)
+        .map(|matched| matched.as_str().eq_ignore_ascii_case("off"))
+}
+
 pub(crate) fn format(tokens: &[Token<'_>], params: &QueryParams, options: FormatOptions) -> String {
     let mut formatter = Formatter::new(tokens, params, options);
     let mut formatted_query = String::new();
+    let mut is_fmt_enabled = true;
+    let mut is_prev_token_fmt_switch = false;
     for (index, token) in tokens.iter().enumerate() {
+        if is_prev_token_fmt_switch {
+            is_prev_token_fmt_switch = false;
+            continue;
+        }
+        if matches!(token.kind, TokenKind::LineComment | TokenKind::BlockComment) {
+            if let Some(is_fmt_off) = check_fmt_off(token.value) {
+                is_fmt_enabled = !is_fmt_off;
+                is_prev_token_fmt_switch = true;
+                continue;
+            }
+        }
         formatter.index = index;
 
+        if !is_fmt_enabled {
+            formatter.format_no_change(token, &mut formatted_query);
+            continue;
+        }
+
         if token.kind == TokenKind::Whitespace {
             // ignore (we do our own whitespace formatting)
         } else if token.kind == TokenKind::LineComment {
@@ -79,16 +108,7 @@ impl<'a> Formatter<'a> {
             self.next_token(1).map_or(false, |current_token| {
                 current_token.kind == TokenKind::Whitespace
                     && self.next_token(2).map_or(false, |next_token| {
-                        matches!(
-                            next_token.kind,
-                            TokenKind::Number
-                                | TokenKind::String
-                                | TokenKind::Word
-                                | TokenKind::ReservedTopLevel
-                                | TokenKind::ReservedTopLevelNoIndent
-                                | TokenKind::ReservedNewline
-                                | TokenKind::Reserved
-                        )
+                        !matches!(next_token.kind, TokenKind::Operator)
                     })
             });
 
@@ -315,4 +335,8 @@ impl<'a> Formatter<'a> {
             None
         }
     }
+
+    fn format_no_change(&self, token: &Token<'_>, query: &mut String) {
+        query.push_str(token.value);
+    }
 }

src/lib.rs

@@ -1653,4 +1653,43 @@ mod tests {
 
         assert_eq!(format(input, &QueryParams::None, options), expected);
     }
+
+    #[test]
+    fn it_recognizes_fmt_off() {
+        let input = indoc!(
+            "SELECT              *     FROM   sometable        
+            WHERE
+            -- comment test here
+                 -- fmt: off
+                first_key.second_key = 1
+                                -- json:first_key.second_key = 1
+                      -- fmt: on
+                AND 
+                   -- fm1t: off
+                first_key.second_key = 1
+                                    --  json:first_key.second_key = 1
+                -- fmt:on"
+        );
+        let options = FormatOptions {
+            indent: Indent::Spaces(4),
+            ..Default::default()
+        };
+        let expected = indoc!(
+            "
+            SELECT
+                *
+            FROM
+                sometable
+            WHERE
+                -- comment test here
+                first_key.second_key = 1
+                                -- json:first_key.second_key = 1
+                AND
+                -- fm1t: off
+                first_key.second_key = 1
+                --  json:first_key.second_key = 1"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, options), expected);
+    }
 }