diff --git a/crates/squawk_fmt/src/fmt.rs b/crates/squawk_fmt/src/fmt.rs index 10f9bab5..d35411d1 100644 --- a/crates/squawk_fmt/src/fmt.rs +++ b/crates/squawk_fmt/src/fmt.rs @@ -125,14 +125,16 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> { } if let Some(group) = &select.group_by_clause() { - doc = doc.append( - Doc::line_or_space() - .append(Doc::text("group by")) - .append(Doc::space()) - .append(Doc::text( - group.group_by_list().unwrap().syntax().to_string(), - )), - ); + let mut group_doc = Doc::line_or_space().append(leading_comments(group.syntax())); + group_doc = group_doc.append(Doc::text("group")).append(Doc::space()); + if let Some(by_token) = group.by_token() { + group_doc = group_doc.append(leading_comments_token(&by_token)); + } + group_doc = group_doc.append(Doc::text("by")).append(Doc::space()); + if let Some(list) = group.group_by_list() { + group_doc = group_doc.append(build_group_by_list(list)); + } + doc = doc.append(group_doc); } doc = doc.append(build_semicolon(select.semicolon_token())); @@ -140,6 +142,10 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> { doc.group() } +fn build_group_by_list<'a>(list: ast::GroupByList) -> Doc<'a> { + leading_comments(list.syntax()).append(Doc::text(list.syntax().to_string())) +} + fn build_semicolon<'a>(semi: Option) -> Doc<'a> { let Some(semi) = semi else { return Doc::nil(); diff --git a/crates/squawk_fmt/tests/after/select.snap b/crates/squawk_fmt/tests/after/select.snap index a2e2cef2..039eec33 100644 --- a/crates/squawk_fmt/tests/after/select.snap +++ b/crates/squawk_fmt/tests/after/select.snap @@ -13,3 +13,5 @@ select select foo as "Quoted Alias" from "Quoted Table"; select 1 as foo; + +select 1 /*a*/ group /* b */ by /*c */ 1; diff --git a/crates/squawk_fmt/tests/before/select.sql b/crates/squawk_fmt/tests/before/select.sql index ec9ba398..1655ee69 100644 --- a/crates/squawk_fmt/tests/before/select.sql +++ b/crates/squawk_fmt/tests/before/select.sql @@ -7,3 +7,5 @@ select 'really long string ' select foo as "Quoted Alias" from "Quoted Table"; select 1 as "foo"; + +select 1 /*a*/group /* b */by/*c */ 1;