From 8b43cbdb3127111a748734dbf9184e4504043d26 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Sun, 23 Aug 2026 15:38:29 -0400 Subject: [PATCH] fmt: support from clause & collation --- crates/squawk_fmt/src/fmt.rs | 275 +++++++++++++++++- crates/squawk_fmt/tests/after/from.snap | 11 + crates/squawk_fmt/tests/after/group_by.snap | 15 + .../squawk_fmt/tests/after/select_expr.snap | 1 + crates/squawk_fmt/tests/before/from.sql | 9 + crates/squawk_fmt/tests/before/group_by.sql | 18 ++ .../squawk_fmt/tests/before/select_expr.sql | 1 + 7 files changed, 316 insertions(+), 14 deletions(-) create mode 100644 crates/squawk_fmt/tests/after/from.snap create mode 100644 crates/squawk_fmt/tests/after/group_by.snap create mode 100644 crates/squawk_fmt/tests/before/from.sql create mode 100644 crates/squawk_fmt/tests/before/group_by.sql diff --git a/crates/squawk_fmt/src/fmt.rs b/crates/squawk_fmt/src/fmt.rs index c6c51960..bfdd5c2b 100644 --- a/crates/squawk_fmt/src/fmt.rs +++ b/crates/squawk_fmt/src/fmt.rs @@ -233,14 +233,11 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> { } } - if let Some(from) = &select.from_clause() { + if let Some(from) = select.from_clause() { doc = doc.append( Doc::line_or_space() - .append(Doc::text("from")) - .append(Doc::space()) - .append(Doc::text( - from.from_items().next().unwrap().syntax().to_string(), - )), + .append(leading_comments(from.syntax())) + .append(build_from_clause(from)), ); } @@ -251,6 +248,15 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> { group_doc = group_doc.append(leading_comments_token(&by_token)); } group_doc = group_doc.append(Doc::text("by")).append(Doc::space()); + if let Some(quantifier) = group.all_or_distinct() { + group_doc = group_doc + .append(leading_comments(quantifier.syntax())) + .append(match quantifier { + ast::AllOrDistinct::All(_) => Doc::text("all"), + ast::AllOrDistinct::Distinct(_) => Doc::text("distinct"), + }) + .append(Doc::space()); + } if let Some(list) = group.group_by_list() { group_doc = group_doc.append(build_group_by_list(list)); } @@ -262,8 +268,240 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> { doc.group() } +fn build_from_clause<'a>(from: ast::FromClause) -> Doc<'a> { + if from.join_exprs().next().is_some() { + todo!("joins are not supported yet") + } + + let from_items: Vec<_> = from + .from_items() + .map(|item| { + let leading = leading_comments(item.syntax()); + let trailing = trailing_comments(item.syntax()); + leading.append(build_from_item(item)).append(trailing) + }) + .collect(); + + Doc::text("from").append(Doc::space()).append( + Doc::list( + Itertools::intersperse( + from_items.into_iter(), + Doc::text(",").append(Doc::line_or_space()), + ) + .collect(), + ) + .nest(2), + ) +} + +fn build_from_item<'a>(item: ast::FromItem) -> Doc<'a> { + match item { + ast::FromItem::RelationFromItem(relation) => build_relation_from_item(relation), + ast::FromItem::FunctionFromItem(_) => { + todo!("function from items are not supported yet") + } + ast::FromItem::ExprFromItem(_) => todo!("expression from items are not supported yet"), + ast::FromItem::ParenFromItem(_) => { + todo!("parenthesized from items are not supported yet") + } + ast::FromItem::RowsFromItem(_) => todo!("rows from items are not supported yet"), + ast::FromItem::GraphTableFromItem(_) => { + todo!("graph_table from items are not supported yet") + } + ast::FromItem::JsonTableFromItem(_) => { + todo!("json_table from items are not supported yet") + } + ast::FromItem::XmlTableFromItem(_) => { + todo!("xmltable from items are not supported yet") + } + } +} + +fn build_relation_from_item<'a>(relation: ast::RelationFromItem) -> Doc<'a> { + let mut doc = if relation.only_token().is_some() { + Doc::text("only").append(Doc::space()) + } else { + Doc::nil() + }; + + if let Some(name) = relation.relation_name_ref() { + doc = doc.append(leading_comments(name.syntax())); + if let Some(path) = name.path_ref() { + doc = doc.append(build_path_ref(&path)); + } + } + if let Some(star) = relation.star_token() { + doc = doc + .append(Doc::space()) + .append(leading_comments_token(&star)) + .append(Doc::text("*")); + } + if let Some(tablesample) = relation.tablesample_clause() { + doc = doc + .append(Doc::space()) + .append(leading_comments(tablesample.syntax())) + .append(build_tablesample_clause(tablesample)); + } + doc.append(build_from_alias(relation.alias())) +} + +fn build_tablesample_clause<'a>(tablesample: ast::TablesampleClause) -> Doc<'a> { + let mut doc = Doc::text("tablesample").append(Doc::space()); + if let Some(call) = tablesample.call_expr() { + doc = doc + .append(leading_comments(call.syntax())) + .append(build_call_expr(call)); + } + if let Some(repeatable) = tablesample.repeatable_clause() { + doc = doc + .append(Doc::space()) + .append(leading_comments(repeatable.syntax())) + .append(Doc::text("repeatable")); + if let Some(l_paren) = repeatable.l_paren_token() { + doc = doc.append(comments_before(l_paren)); + } + doc = doc.append(Doc::text("(")); + if let Some(expr) = repeatable.expr() { + doc = doc + .append(leading_comments(expr.syntax())) + .append(build_expr(expr)); + } + if let Some(r_paren) = repeatable.r_paren_token() { + doc = doc.append(comments_before(r_paren)); + } + doc = doc.append(Doc::text(")")); + } + doc +} + +fn build_from_alias<'a>(alias: Option) -> Doc<'a> { + let Some(alias) = alias else { + return Doc::nil(); + }; + let mut doc = Doc::space().append(leading_comments(alias.syntax())); + if alias.as_token().is_some() { + doc = doc.append(Doc::text("as")).append(Doc::space()); + } + if let Some(name) = alias.name() { + doc = doc + .append(leading_comments(name.syntax())) + .append(build_name(name.syntax())); + } + if alias.columns().is_some() { + todo!("columns in from aliases are not supported yet") + } + doc +} + fn build_group_by_list<'a>(list: ast::GroupByList) -> Doc<'a> { - leading_comments(list.syntax()).append(Doc::text(list.syntax().to_string())) + leading_comments(list.syntax()).append(build_group_bys(list.group_bys())) +} + +fn build_group_bys<'a>(group_bys: impl Iterator) -> Doc<'a> { + Doc::list( + Itertools::intersperse( + group_bys.map(|group_by| { + let leading = leading_comments(group_by.syntax()); + let trailing = trailing_comments(group_by.syntax()); + leading.append(build_group_by(group_by)).append(trailing) + }), + Doc::text(",").append(Doc::line_or_space()), + ) + .collect(), + ) + .nest(2) +} + +fn build_group_by<'a>(group_by: ast::GroupBy) -> Doc<'a> { + match group_by { + ast::GroupBy::GroupingExpr(grouping_expr) => grouping_expr + .expr() + .map(build_expr) + .unwrap_or_else(Doc::nil), + ast::GroupBy::GroupingRollup(rollup) => Doc::text("rollup").append(build_grouping_exprs( + rollup.l_paren_token(), + rollup.exprs(), + rollup.r_paren_token(), + )), + ast::GroupBy::GroupingCube(cube) => Doc::text("cube").append(build_grouping_exprs( + cube.l_paren_token(), + cube.exprs(), + cube.r_paren_token(), + )), + ast::GroupBy::GroupingSets(sets) => { + let mut doc = Doc::text("grouping").append(Doc::space()); + if let Some(sets_token) = sets.sets_token() { + doc = doc.append(leading_comments_token(&sets_token)); + } + doc.append(Doc::text("sets")) + .append(build_grouping_group_bys( + sets.l_paren_token(), + sets.group_bys(), + sets.r_paren_token(), + )) + } + } +} + +fn build_grouping_exprs<'a>( + l_paren: Option, + exprs: impl Iterator, + r_paren: Option, +) -> Doc<'a> { + let exprs: Vec<_> = exprs + .map(|expr| { + let leading = leading_comments(expr.syntax()); + let trailing = trailing_comments(expr.syntax()); + leading.append(build_expr(expr)).append(trailing) + }) + .collect(); + build_grouping_list(l_paren, exprs, r_paren) +} + +fn build_grouping_group_bys<'a>( + l_paren: Option, + group_bys: impl Iterator, + r_paren: Option, +) -> Doc<'a> { + let group_bys = group_bys + .map(|group_by| { + let leading = leading_comments(group_by.syntax()); + let trailing = trailing_comments(group_by.syntax()); + leading.append(build_group_by(group_by)).append(trailing) + }) + .collect(); + build_grouping_list(l_paren, group_bys, r_paren) +} + +fn build_grouping_list<'a>( + l_paren: Option, + items: Vec>, + r_paren: Option, +) -> Doc<'a> { + let mut doc = Doc::nil(); + if let Some(l_paren) = l_paren { + doc = doc.append(comments_before(l_paren)); + } + doc = doc.append(Doc::text("(")); + + if items.is_empty() { + if let Some(r_paren) = r_paren { + doc = doc.append(comments_before(r_paren)); + } + } else { + doc = doc.append( + Doc::list( + Itertools::intersperse( + items.into_iter(), + Doc::text(",").append(Doc::line_or_space()), + ) + .collect(), + ) + .nest(2), + ); + } + + doc.append(Doc::text(")")).group() } fn build_semicolon<'a>(semi: Option) -> Doc<'a> { @@ -743,13 +981,22 @@ fn build_cast_expr<'a>(cast_expr: ast::CastExpr) -> Doc<'a> { } fn build_collate_expr<'a>(collate: ast::Collate) -> Doc<'a> { - build_expr(collate.expr().unwrap()) - .append(Doc::space()) - .append(Doc::text("collate")) - .append(Doc::space()) - .append(Doc::text( - collate.collation_ref().unwrap().syntax().to_string(), - )) + let mut doc = collate.expr().map(build_expr).unwrap_or_else(Doc::nil); + + if let Some(collate_token) = collate.collate_token() { + doc = doc.append(comments_before(collate_token)); + } + doc = doc.append(Doc::space()).append(Doc::text("collate")); + + if let Some(collation) = collate.collation_ref() { + doc = doc + .append(Doc::space()) + .append(leading_comments(collation.syntax())); + if let Some(path) = collation.path_ref() { + doc = doc.append(build_path_ref(&path)); + } + } + doc } fn build_paren_expr<'a>(paren_expr: ast::ParenExpr) -> Doc<'a> { diff --git a/crates/squawk_fmt/tests/after/from.snap b/crates/squawk_fmt/tests/after/from.snap new file mode 100644 index 00000000..38b71d64 --- /dev/null +++ b/crates/squawk_fmt/tests/after/from.snap @@ -0,0 +1,11 @@ +--- +source: crates/squawk_fmt/tests/tests.rs +input_file: crates/squawk_fmt/tests/before/from.sql +--- +select * from foo; +select * from public.foo as f, bar b; +select * from users tablesample bernoulli(10) repeatable(42); +select + * +/* before from */ from /* before item */ only /* before relation */ public /* before dot */./* before table */ foo /* before star */ * /* before alias */ as /* before alias name */ f /* before item comma */, + /* before second item */ other /* before second alias */ o; diff --git a/crates/squawk_fmt/tests/after/group_by.snap b/crates/squawk_fmt/tests/after/group_by.snap new file mode 100644 index 00000000..ec389997 --- /dev/null +++ b/crates/squawk_fmt/tests/after/group_by.snap @@ -0,0 +1,15 @@ +--- +source: crates/squawk_fmt/tests/tests.rs +input_file: crates/squawk_fmt/tests/before/group_by.sql +--- +select 1 group by 1, foo + 2; +select 1 group by all rollup(1, 2), cube(3, 4); +select 1 group by distinct grouping sets((), (1, 2), rollup(3), cube(4)); + +select + 1 +group /* before by */ by /* before distinct */ distinct /* before grouping */ grouping /* before sets */ sets /* before outer paren */(/* before rollup */ rollup /* before rollup paren */(/* before first expression */ 1 /* before expression comma */, + /* before second expression */ 2 /* before rollup close */) /* before group-by comma */, + /* before cube */ cube /* before cube paren */(/* before cube expression */ 3 /* before cube close */) /* before second group-by comma */, + /* before nested grouping */ grouping /* before nested sets */ sets /* before nested paren */(/* before empty tuple */ () /* before nested comma */, + /* before grouping expression */ (4 /* before tuple comma */, /* before tuple expression */ 5 /* before tuple close */) /* before nested close */) /* before outer close */)/* before semicolon */; diff --git a/crates/squawk_fmt/tests/after/select_expr.snap b/crates/squawk_fmt/tests/after/select_expr.snap index 26cd08b2..d690b6b1 100644 --- a/crates/squawk_fmt/tests/after/select_expr.snap +++ b/crates/squawk_fmt/tests/after/select_expr.snap @@ -20,6 +20,7 @@ select ts at time zone 'UTC', 2 ^ 3, name collate "C", + name /* before collate */ collate /* before collation */ public /* before dot */./* before name */ "C", 1 = 1, 1 >= 1, 'foo' ilike 'F%', diff --git a/crates/squawk_fmt/tests/before/from.sql b/crates/squawk_fmt/tests/before/from.sql new file mode 100644 index 00000000..a13b4de4 --- /dev/null +++ b/crates/squawk_fmt/tests/before/from.sql @@ -0,0 +1,9 @@ +select * from foo; +select * from public.foo as f, bar b; +select * from users tablesample bernoulli(10) repeatable (42); +select * +/* before from */ from + /* before item */ only /* before relation */ public /* before dot */ . /* before table */ foo + /* before star */ * + /* before alias */ as /* before alias name */ f /* before item comma */, + /* before second item */ other /* before second alias */ o; diff --git a/crates/squawk_fmt/tests/before/group_by.sql b/crates/squawk_fmt/tests/before/group_by.sql new file mode 100644 index 00000000..bcadeb04 --- /dev/null +++ b/crates/squawk_fmt/tests/before/group_by.sql @@ -0,0 +1,18 @@ +select 1 group by 1, foo + 2; +select 1 group by all rollup (1, 2), cube (3, 4); +select 1 group by distinct grouping sets ((), (1, 2), rollup (3), cube (4)); + +select 1 +group /* before by */ by /* before distinct */ distinct /* before grouping */ grouping /* before sets */ sets /* before outer paren */ ( + /* before rollup */ rollup /* before rollup paren */ ( + /* before first expression */ 1 /* before expression comma */, + /* before second expression */ 2 /* before rollup close */ + ) /* before group-by comma */, + /* before cube */ cube /* before cube paren */ ( + /* before cube expression */ 3 /* before cube close */ + ) /* before second group-by comma */, + /* before nested grouping */ grouping /* before nested sets */ sets /* before nested paren */ ( + /* before empty tuple */ () /* before nested comma */, + /* before grouping expression */ (4 /* before tuple comma */, /* before tuple expression */ 5 /* before tuple close */) /* before nested close */ + ) /* before outer close */ +) /* before semicolon */; diff --git a/crates/squawk_fmt/tests/before/select_expr.sql b/crates/squawk_fmt/tests/before/select_expr.sql index 1161967f..96ee9506 100644 --- a/crates/squawk_fmt/tests/before/select_expr.sql +++ b/crates/squawk_fmt/tests/before/select_expr.sql @@ -16,6 +16,7 @@ select ts at time zone 'UTC', 2 ^ 3, name collate "C", + name /* before collate */ collate /* before collation */ public /* before dot */ . /* before name */ "C", 1 = 1, 1 >= 1, 'foo' ilike 'F%',