feat: Uniform the behaviour of block and top-level span options

lu-zero · shssoichiro · commit c781cb949735 · 2025-01-19T23:45:54.000-05:00
And fix the span computation.
diff --git a/src/formatter.rs b/src/formatter.rs
@@ -148,7 +148,8 @@ struct Formatter<'a> {
     options: &'a FormatOptions<'a>,
     indentation: Indentation<'a>,
     inline_block: InlineBlock,
-    line_start: usize,
+    span_len: usize,
+    block_level: usize,
 }
 
 impl<'a> Formatter<'a> {
@@ -165,7 +166,8 @@ impl<'a> Formatter<'a> {
                 options.max_inline_block,
                 options.max_inline_arguments.is_none(),
             ),
-            line_start: 0,
+            span_len: 0,
+            block_level: 0,
         }
     }
 
@@ -209,11 +211,11 @@ impl<'a> Formatter<'a> {
         self.add_new_line(query);
         self.indentation.increase_top_level();
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
-        let len = self.top_level_tokens_span();
+        self.span_len = self.top_level_tokens_span();
         if self
             .options
             .max_inline_top_level
-            .map_or(true, |limit| limit < len)
+            .map_or(true, |limit| limit < self.span_len)
         {
             self.add_new_line(query);
         } else {
@@ -232,7 +234,7 @@ impl<'a> Formatter<'a> {
         if self
             .options
             .max_inline_arguments
-            .map_or(true, |limit| limit < self.line_len_next(query))
+            .map_or(true, |limit| limit < self.span_len)
         {
             self.add_new_line(query);
         } else {
@@ -256,6 +258,7 @@ impl<'a> Formatter<'a> {
 
     // Opening parentheses increase the block indent level and start a new line
     fn format_opening_parentheses(&mut self, token: &Token<'_>, query: &mut String) {
+        self.block_level += 1;
         const PRESERVE_WHITESPACE_FOR: &[TokenKind] = &[
             TokenKind::Whitespace,
             TokenKind::OpenParen,
@@ -304,6 +307,7 @@ impl<'a> Formatter<'a> {
 
     // Closing parentheses decrease the block indent level
     fn format_closing_parentheses(&mut self, token: &Token<'_>, query: &mut String) {
+        self.block_level = self.block_level.saturating_sub(1);
         let mut token = token.clone();
         let value = match (
             self.options.uppercase,
@@ -367,7 +371,7 @@ impl<'a> Formatter<'a> {
         }
 
         if matches!((self.previous_top_level_reserved_word, self.options.max_inline_arguments),
-            (Some(word), Some(limit)) if word.value.to_lowercase() == "select" && limit > self.line_len_next(query))
+            (Some(word), Some(limit)) if ["select", "from"].contains(&word.value.to_lowercase().as_str()) && limit > self.span_len)
         {
             return;
         }
@@ -394,7 +398,7 @@ impl<'a> Formatter<'a> {
         }
     }
 
-    fn add_new_line(&mut self, query: &mut String) {
+    fn add_new_line(&self, query: &mut String) {
         self.trim_spaces_end(query);
         if self.options.inline {
             query.push(' ');
@@ -403,7 +407,6 @@ impl<'a> Formatter<'a> {
         if !query.ends_with('\n') {
             query.push('\n');
         }
-        self.line_start = query.len();
         query.push_str(&self.indentation.get_indent());
     }
 
@@ -490,16 +493,25 @@ impl<'a> Formatter<'a> {
         }
     }
 
-    fn line_len_next(&self, query: &str) -> usize {
-        query.len() - self.line_start + self.next_token(1).map_or(0, |t| t.value.len())
-    }
-
     fn top_level_tokens_span(&self) -> usize {
         assert_eq!(self.tokens[self.index].kind, TokenKind::ReservedTopLevel);
+        let mut block_level = self.block_level;
 
         self.tokens[self.index..]
             .iter()
             .skip(1)
+            .take_while(|token| match token.kind {
+                TokenKind::OpenParen => {
+                    block_level += 1;
+                    true
+                }
+                TokenKind::CloseParen => {
+                    block_level = block_level.saturating_sub(1);
+                    block_level > self.block_level
+                }
+                TokenKind::ReservedTopLevel => false,
+                _ => true,
+            })
             .map(|token| token.value.len())
             .sum()
     }
diff --git a/src/inline_block.rs b/src/inline_block.rs
@@ -81,6 +81,6 @@ impl InlineBlock {
             } else {
                 false
             }
-            || token.value.to_lowercase() == "case"
+            || ["case", "end"].contains(&token.value.to_lowercase().as_str())
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -197,7 +197,7 @@ mod tests {
     }
 
     #[test]
-    fn split_select_arguments_over_lines() {
+    fn split_select_arguments_inline_top_level() {
         let input = indoc! {
             "
             SELECT
@@ -212,16 +212,14 @@ mod tests {
             FROM foo;"
         };
         let options = FormatOptions {
-            max_inline_arguments: Some(20),
+            max_inline_arguments: Some(50),
+            max_inline_top_level: Some(50),
             ..Default::default()
         };
         let expected = indoc! {
             "
-            SELECT
-              a, b, c, d, e,
-              f, g, h
-            FROM
-              foo;"
+            SELECT a, b, c, d, e, f, g, h
+            FROM foo;"
         };
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
@@ -1514,6 +1512,31 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn it_formats_case_when_inside_select_inlined_top_level() {
+        let input =
+            "SELECT foo, bar, CASE baz WHEN 'one' THEN 1 WHEN 'two' THEN 2 ELSE 3 END FROM table";
+        let options = FormatOptions {
+            max_inline_top_level: Some(50),
+            ..Default::default()
+        };
+        let expected = indoc!(
+            "
+            SELECT
+              foo,
+              bar,
+              CASE
+                baz
+                WHEN 'one' THEN 1
+                WHEN 'two' THEN 2
+                ELSE 3
+              END
+            FROM table"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_formats_case_when_with_an_expression() {
         let input = "CASE toString(getNumber()) WHEN 'one' THEN 1 WHEN 'two' THEN 2 WHEN 'three' THEN 3 ELSE 4 END;";
@@ -1958,6 +1981,33 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn format_nested_select() {
+        let input = "WITH a AS ( SELECT a, b, c FROM t WHERE a > 100 ), aa AS ( SELECT field FROM table ) SELECT b, field FROM a, aa;";
+        let options = FormatOptions {
+            max_inline_arguments: Some(10),
+            max_inline_top_level: Some(9),
+            ..Default::default()
+        };
+        let expected = indoc! {
+            "
+            WITH a AS (
+              SELECT a, b, c
+              FROM t
+              WHERE a > 100
+            ),
+            aa AS (
+              SELECT field
+              FROM table
+            )
+            SELECT
+              b,
+              field
+            FROM a, aa;"
+        };
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_converts_keywords_nothing_when_no_option_passed_in() {
         let input = "select distinct * frOM foo left join bar WHERe cola > 1 and colb = 3";

Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,6 @@ impl InlineBlock {`
`81`	`81`	`} else {`
`82`	`82`	`false`
`83`	`83`	`}`
`84`		`- \|\| token.value.to_lowercase() == "case"`
	`84`	`+ \|\| ["case", "end"].contains(&token.value.to_lowercase().as_str())`
`85`	`85`	`}`
`86`	`86`	`}`