Correctly format array index notation

Author: lu-zero

src/formatter.rs

@@ -369,7 +369,8 @@ impl<'a> Formatter<'a> {
         ];
 
         const ADD_WHITESPACE_BETWEEN: &[TokenKind] = &[TokenKind::CloseParen, TokenKind::Reserved];
-
+        const BEFORE_ARRAY: &[TokenKind] =
+            &[TokenKind::CloseParen, TokenKind::Word, TokenKind::Reserved];
         let inlined = self.inline_block.begin_if_possible(self.tokens, self.index);
         let previous_non_whitespace_token = self.previous_non_whitespace_token(1);
         let fold_in_top_level = !inlined
@@ -387,13 +388,16 @@ impl<'a> Formatter<'a> {
         // Take out the preceding space unless there was whitespace there in the original query
         // or another opening parens or line comment
         let previous_token = self.previous_token(1);
-        if previous_token.is_none()
-            || !PRESERVE_WHITESPACE_FOR.contains(&previous_token.unwrap().kind)
+        if previous_token.is_none_or(|t| !PRESERVE_WHITESPACE_FOR.contains(&t.kind))
+            || previous_non_whitespace_token
+                .is_some_and(|t| token.value == "[" && BEFORE_ARRAY.contains(&t.kind))
         {
             self.trim_spaces_end(query);
         }
 
-        if previous_non_whitespace_token.is_some_and(|t| ADD_WHITESPACE_BETWEEN.contains(&t.kind)) {
+        if previous_non_whitespace_token
+            .is_some_and(|t| token.value != "[" && ADD_WHITESPACE_BETWEEN.contains(&t.kind))
+        {
             self.trim_spaces_end(query);
             query.push(' ');
         }

src/lib.rs

@@ -554,6 +554,21 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn it_formats_array_index_notation() {
+        let input = "SELECT a [ 1 ] + b [ 2 ] [   5+1 ] > c [3] ;";
+        let options = FormatOptions {
+            dialect: Dialect::PostgreSql,
+            ..Default::default()
+        };
+        let expected = indoc!(
+            "
+            SELECT
+              a[1] + b[2][5 + 1] > c[3];"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
     #[test]
     fn it_formats_limit_of_single_value_and_offset() {
         let input = "LIMIT 5 OFFSET 8;";