Skip to content

Commit 02e0af9

Browse files
lu-zerolu-zero
committed
feat: Support complex insert queries
And also provide better nested inline support.
1 parent ff290d6 commit 02e0af9

File tree

4 files changed

+164
-30
lines changed

4 files changed

+164
-30
lines changed

‎src/formatter.rs‎

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -89,11 +89,15 @@ pub(crate) fn format(
8989
}
9090
TokenKind::ReservedTopLevel => {
9191
formatter.format_top_level_reserved_word(token, &mut formatted_query);
92-
formatter.indentation.set_previous_top_level(token);
92+
formatter.set_top_level_span(token);
9393
}
9494
TokenKind::ReservedTopLevelNoIndent => {
9595
formatter.format_top_level_reserved_word_no_indent(token, &mut formatted_query);
96-
formatter.indentation.set_previous_top_level(token);
96+
formatter.set_top_level_span(token);
97+
}
98+
TokenKind::ReservedNewlineAfter => {
99+
formatter.format_newline_after_reserved_word(token, &mut formatted_query);
100+
formatter.set_top_level_span(token);
97101
}
98102
TokenKind::ReservedNewline => {
99103
formatter.format_newline_reserved_word(token, &mut formatted_query);
@@ -164,6 +168,11 @@ impl<'a> Formatter<'a> {
164168
}
165169
}
166170

171+
fn set_top_level_span(&mut self, token: &'a Token<'a>) {
172+
let span_info = self.top_level_tokens_info();
173+
self.indentation.set_previous_top_level(token, span_info);
174+
}
175+
167176
fn format_line_comment(&mut self, token: &Token<'_>, query: &mut String) {
168177
let is_whitespace_followed_by_special_token =
169178
self.next_token(1).map_or(false, |current_token| {
@@ -249,6 +258,31 @@ impl<'a> Formatter<'a> {
249258
}
250259
}
251260

261+
fn format_newline_after_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
262+
let span_info = self.top_level_tokens_info();
263+
264+
let newline_before = match (
265+
self.options.max_inline_top_level,
266+
self.indentation.previous_top_level_reserved(),
267+
) {
268+
(Some(limit), Some((_, span))) => limit < span.full_span + token.value.len(),
269+
_ => true,
270+
};
271+
272+
if newline_before {
273+
self.indentation.decrease_top_level();
274+
self.add_new_line(query);
275+
} else {
276+
self.trim_spaces_end(query);
277+
query.push(' ');
278+
}
279+
280+
query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
281+
282+
self.indentation.increase_top_level(span_info);
283+
self.add_new_line(query);
284+
}
285+
252286
fn format_newline_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
253287
if !self.inline_block.is_active()
254288
&& self
@@ -395,12 +429,13 @@ impl<'a> Formatter<'a> {
395429
return;
396430
}
397431

398-
if matches!((self.indentation.previous_top_level_reserved(), self.options.max_inline_arguments),
399-
(Some(word), Some(limit)) if ["select", "from"].contains(&word.value.to_lowercase().as_str()) &&
400-
limit > self.indentation.span())
401-
{
402-
return;
432+
if let Some((_, span)) = self.indentation.previous_top_level_reserved() {
433+
let limit = self.options.max_inline_arguments.unwrap_or(0);
434+
if limit > span.full_span {
435+
return;
436+
}
403437
}
438+
404439
self.add_new_line(query);
405440
}
406441

@@ -534,11 +569,18 @@ impl<'a> Formatter<'a> {
534569
break;
535570
}
536571
}
537-
TokenKind::ReservedTopLevel | TokenKind::ReservedTopLevelNoIndent => {
572+
TokenKind::ReservedTopLevel
573+
| TokenKind::ReservedTopLevelNoIndent
574+
| TokenKind::ReservedNewlineAfter => {
538575
if block_level == self.block_level {
539576
break;
540577
}
541578
}
579+
TokenKind::Whitespace => {
580+
full_span += 1;
581+
continue;
582+
}
583+
542584
_ => {}
543585
}
544586

‎src/indentation.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::{tokenizer::Token, FormatOptions, Indent, SpanInfo};
22

33
#[derive(Debug, Default)]
44
struct PreviousTokens<'a> {
5-
top_level_reserved: Option<&'a Token<'a>>,
5+
top_level_reserved: Option<(&'a Token<'a>, SpanInfo)>,
66
reserved: Option<&'a Token<'a>>,
77
}
88

@@ -83,12 +83,12 @@ impl<'a> Indentation<'a> {
8383
}
8484
}
8585

86-
pub fn set_previous_top_level(&mut self, token: &'a Token<'a>) {
86+
pub fn set_previous_top_level(&mut self, token: &'a Token<'a>, span_info: SpanInfo) {
8787
if let Some(previous) = self.previous.last_mut() {
88-
previous.top_level_reserved = Some(token);
88+
previous.top_level_reserved = Some((token, span_info));
8989
} else {
9090
self.previous.push(PreviousTokens {
91-
top_level_reserved: Some(token),
91+
top_level_reserved: Some((token, span_info)),
9292
reserved: Some(token),
9393
});
9494
}
@@ -106,13 +106,13 @@ impl<'a> Indentation<'a> {
106106
}
107107
}
108108

109-
pub fn previous_top_level_reserved(&'a self) -> Option<&'a Token<'a>> {
109+
pub fn previous_top_level_reserved(&'a self) -> Option<(&'a Token<'a>, &'a SpanInfo)> {
110110
if let Some(PreviousTokens {
111111
top_level_reserved,
112112
reserved: _,
113113
}) = self.previous.last()
114114
{
115-
top_level_reserved.as_deref()
115+
top_level_reserved.as_ref().map(|&(t, ref s)| (t, s))
116116
} else {
117117
None
118118
}

‎src/lib.rs‎

Lines changed: 61 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ pub enum QueryParams {
9797
None,
9898
}
9999

100-
#[derive(Default)]
100+
#[derive(Default, Debug)]
101101
pub(crate) struct SpanInfo {
102102
pub full_span: usize,
103103
// potentially comma span info here
@@ -635,6 +635,40 @@ mod tests {
635635
assert_eq!(format(input, &QueryParams::None, &options), expected);
636636
}
637637

638+
#[test]
639+
fn it_formats_complex_insert_query() {
640+
let input = "
641+
INSERT INTO t(id, a, min, max) SELECT input.id, input.a, input.min, input.max FROM ( SELECT id, a, min, max FROM foo WHERE a IN ('a', 'b') ) AS input WHERE (SELECT true FROM condition) ON CONFLICT ON CONSTRAINT a_id_key DO UPDATE SET id = EXCLUDED.id, a = EXCLUDED.severity, min = EXCLUDED.min, max = EXCLUDED.max RETURNING *; ";
642+
let max_line = 50;
643+
let options = FormatOptions {
644+
max_inline_block: max_line,
645+
max_inline_arguments: Some(max_line),
646+
max_inline_top_level: Some(max_line),
647+
..Default::default()
648+
};
649+
650+
let expected = indoc!(
651+
"
652+
INSERT INTO t(id, a, min, max)
653+
SELECT input.id, input.a, input.min, input.max
654+
FROM
655+
(
656+
SELECT id, a, min, max
657+
FROM foo
658+
WHERE a IN ('a', 'b')
659+
) AS input
660+
WHERE (SELECT true FROM condition)
661+
ON CONFLICT ON CONSTRAINT a_id_key DO UPDATE SET
662+
id = EXCLUDED.id,
663+
a = EXCLUDED.severity,
664+
min = EXCLUDED.min,
665+
max = EXCLUDED.max
666+
RETURNING *;"
667+
);
668+
669+
assert_eq!(format(input, &QueryParams::None, &options), expected);
670+
}
671+
638672
#[test]
639673
fn it_keeps_short_parenthesized_list_with_nested_parenthesis_on_single_line() {
640674
let input = "SELECT (a + b * (c - NOW()));";
@@ -731,6 +765,26 @@ mod tests {
731765
assert_eq!(format(input, &QueryParams::None, &options), expected);
732766
}
733767

768+
#[test]
769+
fn it_formats_simple_update_query_inlining_set() {
770+
let input = "UPDATE Customers SET ContactName='Alfred Schmidt', City='Hamburg' WHERE CustomerName='Alfreds Futterkiste';";
771+
let options = FormatOptions {
772+
max_inline_top_level: Some(20),
773+
max_inline_arguments: Some(10),
774+
..Default::default()
775+
};
776+
let expected = indoc!(
777+
"
778+
UPDATE Customers SET
779+
ContactName = 'Alfred Schmidt',
780+
City = 'Hamburg'
781+
WHERE
782+
CustomerName = 'Alfreds Futterkiste';"
783+
);
784+
785+
assert_eq!(format(input, &QueryParams::None, &options), expected);
786+
}
787+
734788
#[test]
735789
fn it_formats_simple_delete_query() {
736790
let input = "DELETE FROM Customers WHERE CustomerName='Alfred' AND Phone=5002132;";
@@ -1644,13 +1698,10 @@ mod tests {
16441698
#[test]
16451699
fn it_formats_case_when_inside_an_order_by() {
16461700
let input = "SELECT a, created_at FROM b ORDER BY (CASE $3 WHEN 'created_at_asc' THEN created_at END) ASC, (CASE $3 WHEN 'created_at_desc' THEN created_at END) DESC;";
1647-
let max_line = 120;
1701+
let max_line = 80;
16481702
let options = FormatOptions {
16491703
max_inline_block: max_line,
16501704
max_inline_arguments: Some(max_line),
1651-
joins_as_top_level: true,
1652-
uppercase: Some(true),
1653-
ignore_case_convert: Some(vec!["status"]),
16541705
..Default::default()
16551706
};
16561707

@@ -2314,7 +2365,10 @@ from
23142365

23152366
#[test]
23162367
fn it_formats_blocks_inline_or_not() {
2317-
let input = " UPDATE t SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2
2368+
let input = " UPDATE t
2369+
2370+
2371+
SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2
23182372
THEN NULL ELSE COALESCE($3, b) END, b = CASE WHEN $4 THEN NULL ELSE
23192373
COALESCE($5, b) END, s = (SELECT true FROM bar WHERE bar.foo = $99 AND bar.foo > $100),
23202374
c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END,
@@ -2328,8 +2382,7 @@ from
23282382
};
23292383
let expected = indoc!(
23302384
"
2331-
UPDATE t
2332-
SET
2385+
UPDATE t SET
23332386
o = ($5 + $6 + $7 + $8),
23342387
a = CASE WHEN $2 THEN NULL ELSE COALESCE($3, b) END,
23352388
b = CASE WHEN $4 THEN NULL ELSE COALESCE($5, b) END,

‎src/tokenizer.rs‎

Lines changed: 47 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ pub(crate) enum TokenKind {
7575
ReservedTopLevel,
7676
ReservedTopLevelNoIndent,
7777
ReservedNewline,
78+
ReservedNewlineAfter,
7879
Operator,
7980
OpenParen,
8081
CloseParen,
@@ -393,6 +394,7 @@ fn get_reserved_word_token<'a>(
393394

394395
alt((
395396
get_top_level_reserved_token(last_reserved_top_level_token),
397+
get_newline_after_reserved_token(),
396398
get_newline_reserved_token(last_reserved_token),
397399
get_join_token(),
398400
get_top_level_reserved_token_no_indent,
@@ -499,14 +501,21 @@ fn get_top_level_reserved_token<'a>(
499501
input.to_ascii_uppercase().find(final_word).unwrap_or(0) + final_word.len();
500502
let token = input.next_slice(input_end_pos);
501503

502-
let kind = if token == "EXCEPT"
503-
&& last_reserved_top_level_token.is_some()
504-
&& last_reserved_top_level_token.as_ref().unwrap().value == "SELECT"
505-
{
504+
let kind = match token {
505+
"EXCEPT"
506+
if last_reserved_top_level_token.is_some()
507+
&& last_reserved_top_level_token.as_ref().unwrap().value == "SELECT" =>
506508
// If the query state doesn't allow EXCEPT, treat it as a regular word
507-
TokenKind::Word
508-
} else {
509-
TokenKind::ReservedTopLevel
509+
{
510+
TokenKind::Word
511+
}
512+
"SET"
513+
if last_reserved_top_level_token.is_some()
514+
&& last_reserved_top_level_token.as_ref().unwrap().value == "UPDATE" =>
515+
{
516+
TokenKind::ReservedNewlineAfter
517+
}
518+
_ => TokenKind::ReservedTopLevel,
510519
};
511520

512521
Ok(Token {
@@ -583,6 +592,35 @@ fn get_join_token<'a>() -> impl Parser<&'a str, Token<'a>, ContextError> {
583592
}
584593
}
585594

595+
fn get_newline_after_reserved_token<'a>() -> impl Parser<&'a str, Token<'a>, ContextError> {
596+
move |input: &mut &'a str| {
597+
let uc_input: String = get_uc_words(input, 3);
598+
let mut uc_input = uc_input.as_str();
599+
600+
let mut on_conflict = alt((
601+
terminated("DO NOTHING", end_of_word),
602+
terminated("DO UPDATE SET", end_of_word),
603+
));
604+
605+
let result: Result<&str> = on_conflict.parse_next(&mut uc_input);
606+
607+
if let Ok(token) = result {
608+
let final_word = token.split(' ').next_back().unwrap();
609+
let input_end_pos =
610+
input.to_ascii_uppercase().find(final_word).unwrap() + final_word.len();
611+
let token = input.next_slice(input_end_pos);
612+
let kind = TokenKind::ReservedNewlineAfter;
613+
Ok(Token {
614+
kind,
615+
value: token,
616+
key: None,
617+
})
618+
} else {
619+
Err(ParserError::from_input(input))
620+
}
621+
}
622+
}
623+
586624
fn get_newline_reserved_token<'a>(
587625
last_reserved_token: Option<Token<'a>>,
588626
) -> impl Parser<&'a str, Token<'a>, ContextError> {
@@ -1070,11 +1108,12 @@ fn get_plain_reserved_one_token<'i>(input: &mut &'i str) -> Result<Token<'i>> {
10701108
}
10711109

10721110
fn get_plain_reserved_two_token<'i>(input: &mut &'i str) -> Result<Token<'i>> {
1073-
let uc_input = get_uc_words(input, 2);
1111+
let uc_input = get_uc_words(input, 3);
10741112
let mut uc_input = uc_input.as_str();
10751113
let result: Result<&str> = alt((
10761114
terminated("CHARACTER SET", end_of_word),
10771115
terminated("ON CONFLICT", end_of_word),
1116+
terminated("ON CONSTRAINT", end_of_word),
10781117
terminated("ON DELETE", end_of_word),
10791118
terminated("ON UPDATE", end_of_word),
10801119
terminated("DISTINCT FROM", end_of_word),

0 commit comments

Comments
 (0)