fix #205: support within group and filter keywords

tconbeer · Aug 2, 2022 · bcf660d · bcf660d
1 parent e2c30ac
commit bcf660d
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 
 -   sqlfmt now standardizes whitespace inside word tokens ([#201](https://github.com/tconbeer/sqlfmt/issues/201))
 -   `using` is now treated as a word operator. It gets a space before its brackets and merging with surrounding lines is now much improved ([#218](https://github.com/tconbeer/sqlfmt/issues/218) - thank you [@nfcampos](https://github.com/nfcampos)!)
+-   `within group` and `filter` are now treated like `over`, and the formatting of those aggregate clauses is improved ([#205](https://github.com/tconbeer/sqlfmt/issues/205))
 
 ## [0.10.0] - 2022-08-02
 

diff --git a/src/sqlfmt/dialect.py b/src/sqlfmt/dialect.py
@@ -296,9 +296,10 @@ def __init__(self) -> None:
                     ),
                 ),
                 Rule(
-                    name="over",
+                    name="agg_modifiers",
                     priority=922,
-                    pattern=group(r"over") + group(r"\W", r"$"),
+                    pattern=group(r"over", r"within\s+group", r"filter")
+                    + group(r"\W", r"$"),
                     action=partial(
                         actions.add_node_to_buffer,
                         token_type=TokenType.TIGHT_WORD_OPERATOR,

diff --git a/tests/data/unformatted/103_window_functions.sql b/tests/data/unformatted/103_window_functions.sql
@@ -3,7 +3,8 @@ select
     sum(a) over () as b,
     row_number() over () as c,
     count(case when a is null then 1 end) over (partition by user_id, date_trunc('year', performed_at)) as d,
-    first_value(a ignore nulls) over (partition by user_id order by performed_at desc rows between unbounded preceding and unbounded following) as e
+    first_value(a ignore nulls) over (partition by user_id order by performed_at desc rows between unbounded preceding and unbounded following) as e,
+    count(*) filter (WHERE a is null) over (partition by user_id, date_trunc('year', performed_at)) as f
 from
     my_table
 )))))__SQLFMT_OUTPUT__(((((
@@ -18,5 +19,8 @@ select
         partition by user_id
         order by performed_at desc
         rows between unbounded preceding and unbounded following
-    ) as e
+    ) as e,
+    count(*) filter (where a is null) over (
+        partition by user_id, date_trunc('year', performed_at)
+    ) as f
 from my_table
diff --git a/tests/data/unformatted/118_within_group.sql b/tests/data/unformatted/118_within_group.sql
@@ -0,0 +1,26 @@
+select
+    ARRAY_AGG(DISTINCT foobar.my_special_foo) WITHIN GROUP (ORDER BY foobar.another_foo)    AS a_very_long_alias,
+    ARRAY_AGG(DISTINCT foobar.my_special_foo || foobar.another_very_long_foo) WITHIN GROUP (ORDER BY foobar.another_foo desc)    AS a_very_long_alias,
+    ARRAY_AGG(DISTINCT foobar.my_special_foo) WITHIN GROUP (ORDER BY foobar.another_foo) FILTER (WHERE barbar = "bazbaz" and bazbaz = "quxqux")   AS something_else,
+    ARRAY_AGG(DISTINCT foobar.my_special_foo) WITHIN GROUP (ORDER BY foobar.another_foo) FILTER (WHERE barbar = "bazbaz" and bazbaz = "quxqux" and something_else_quite_long = "a long literal")   AS something_else
+from my_table as foobar
+)))))__SQLFMT_OUTPUT__(((((
+select
+    array_agg(distinct foobar.my_special_foo) within group (
+        order by foobar.another_foo
+    ) as a_very_long_alias,
+    array_agg(
+        distinct foobar.my_special_foo || foobar.another_very_long_foo
+    ) within group (order by foobar.another_foo desc) as a_very_long_alias,
+    array_agg(distinct foobar.my_special_foo) within group (
+        order by foobar.another_foo
+    ) filter (where barbar = "bazbaz" and bazbaz = "quxqux") as something_else,
+    array_agg(distinct foobar.my_special_foo) within group (
+        order by foobar.another_foo
+    ) filter (
+        where
+            barbar = "bazbaz"
+            and bazbaz = "quxqux"
+            and something_else_quite_long = "a long literal"
+    ) as something_else
+from my_table as foobar
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
@@ -29,6 +29,8 @@
         "unformatted/114_unions.sql",
         "unformatted/115_select_star_except.sql",
         "unformatted/116_chained_booleans.sql",
+        "unformatted/117_whitespace_in_tokens.sql",
+        "unformatted/118_within_group.sql",
         "unformatted/200_base_model.sql",
         "unformatted/201_basic_snapshot.sql",
         "unformatted/202_unpivot_macro.sql",

diff --git a/tests/unit_tests/test_dialect.py b/tests/unit_tests/test_dialect.py
@@ -156,7 +156,10 @@ def test_rule_props_are_unique(self, dialect: Polyglot) -> None:
             ("main", "word_operator", "grouping sets"),
             ("main", "word_operator", "cube"),
             ("main", "word_operator", "rollup"),
-            ("main", "over", "over"),
+            ("main", "word_operator", "using"),
+            ("main", "agg_modifiers", "over"),
+            ("main", "agg_modifiers", "within group"),
+            ("main", "agg_modifiers", "filter"),
             ("main", "as", "as"),
             ("main", "on", "on"),
             ("main", "boolean_operator", "AND"),