Skip to content

Commit 2ab0257

Browse files
lu-zerolu-zero
committed
feat: Support complex insert queries
And also provide better nested inline support.
1 parent 6f5df13 commit 2ab0257

File tree

4 files changed

+163
-29
lines changed

4 files changed

+163
-29
lines changed

‎src/formatter.rs‎

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,15 @@ pub(crate) fn format(
8888
}
8989
TokenKind::ReservedTopLevel => {
9090
formatter.format_top_level_reserved_word(token, &mut formatted_query);
91-
formatter.indentation.set_previous_top_level(token);
91+
formatter.set_top_level_span(token);
9292
}
9393
TokenKind::ReservedTopLevelNoIndent => {
9494
formatter.format_top_level_reserved_word_no_indent(token, &mut formatted_query);
95-
formatter.indentation.set_previous_top_level(token);
95+
formatter.set_top_level_span(token);
96+
}
97+
TokenKind::ReservedNewlineAfter => {
98+
formatter.format_newline_after_reserved_word(token, &mut formatted_query);
99+
formatter.set_top_level_span(token);
96100
}
97101
TokenKind::ReservedNewline => {
98102
formatter.format_newline_reserved_word(token, &mut formatted_query);
@@ -163,6 +167,11 @@ impl<'a> Formatter<'a> {
163167
}
164168
}
165169

170+
fn set_top_level_span(&mut self, token: &'a Token<'a>) {
171+
let span_info = self.top_level_tokens_info();
172+
self.indentation.set_previous_top_level(token, span_info);
173+
}
174+
166175
fn format_line_comment(&mut self, token: &Token<'_>, query: &mut String) {
167176
let is_whitespace_followed_by_special_token =
168177
self.next_token(1).is_some_and(|current_token| {
@@ -248,6 +257,31 @@ impl<'a> Formatter<'a> {
248257
}
249258
}
250259

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

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

@@ -533,11 +568,18 @@ impl<'a> Formatter<'a> {
533568
break;
534569
}
535570
}
536-
TokenKind::ReservedTopLevel | TokenKind::ReservedTopLevelNoIndent => {
571+
TokenKind::ReservedTopLevel
572+
| TokenKind::ReservedTopLevelNoIndent
573+
| TokenKind::ReservedNewlineAfter => {
537574
if block_level == self.block_level {
538575
break;
539576
}
540577
}
578+
TokenKind::Whitespace => {
579+
full_span += 1;
580+
continue;
581+
}
582+
541583
_ => {}
542584
}
543585

‎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
@@ -672,6 +672,40 @@ mod tests {
672672
assert_eq!(format(input, &QueryParams::None, &options), expected);
673673
}
674674

675+
#[test]
676+
fn it_formats_complex_insert_query() {
677+
let input = "
678+
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 *; ";
679+
let max_line = 50;
680+
let options = FormatOptions {
681+
max_inline_block: max_line,
682+
max_inline_arguments: Some(max_line),
683+
max_inline_top_level: Some(max_line),
684+
..Default::default()
685+
};
686+
687+
let expected = indoc!(
688+
"
689+
INSERT INTO t(id, a, min, max)
690+
SELECT input.id, input.a, input.min, input.max
691+
FROM
692+
(
693+
SELECT id, a, min, max
694+
FROM foo
695+
WHERE a IN ('a', 'b')
696+
) AS input
697+
WHERE (SELECT true FROM condition)
698+
ON CONFLICT ON CONSTRAINT a_id_key DO UPDATE SET
699+
id = EXCLUDED.id,
700+
a = EXCLUDED.severity,
701+
min = EXCLUDED.min,
702+
max = EXCLUDED.max
703+
RETURNING *;"
704+
);
705+
706+
assert_eq!(format(input, &QueryParams::None, &options), expected);
707+
}
708+
675709
#[test]
676710
fn it_keeps_short_parenthesized_list_with_nested_parenthesis_on_single_line() {
677711
let input = "SELECT (a + b * (c - NOW()));";
@@ -768,6 +802,26 @@ mod tests {
768802
assert_eq!(format(input, &QueryParams::None, &options), expected);
769803
}
770804

805+
#[test]
806+
fn it_formats_simple_update_query_inlining_set() {
807+
let input = "UPDATE Customers SET ContactName='Alfred Schmidt', City='Hamburg' WHERE CustomerName='Alfreds Futterkiste';";
808+
let options = FormatOptions {
809+
max_inline_top_level: Some(20),
810+
max_inline_arguments: Some(10),
811+
..Default::default()
812+
};
813+
let expected = indoc!(
814+
"
815+
UPDATE Customers SET
816+
ContactName = 'Alfred Schmidt',
817+
City = 'Hamburg'
818+
WHERE
819+
CustomerName = 'Alfreds Futterkiste';"
820+
);
821+
822+
assert_eq!(format(input, &QueryParams::None, &options), expected);
823+
}
824+
771825
#[test]
772826
fn it_formats_simple_delete_query() {
773827
let input = "DELETE FROM Customers WHERE CustomerName='Alfred' AND Phone=5002132;";
@@ -1681,13 +1735,10 @@ mod tests {
16811735
#[test]
16821736
fn it_formats_case_when_inside_an_order_by() {
16831737
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;";
1684-
let max_line = 120;
1738+
let max_line = 80;
16851739
let options = FormatOptions {
16861740
max_inline_block: max_line,
16871741
max_inline_arguments: Some(max_line),
1688-
joins_as_top_level: true,
1689-
uppercase: Some(true),
1690-
ignore_case_convert: Some(vec!["status"]),
16911742
..Default::default()
16921743
};
16931744

@@ -2351,7 +2402,10 @@ from
23512402

23522403
#[test]
23532404
fn it_formats_blocks_inline_or_not() {
2354-
let input = " UPDATE t SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2
2405+
let input = " UPDATE t
2406+
2407+
2408+
SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2
23552409
THEN NULL ELSE COALESCE($3, b) END, b = CASE WHEN $4 THEN NULL ELSE
23562410
COALESCE($5, b) END, s = (SELECT true FROM bar WHERE bar.foo = $99 AND bar.foo > $100),
23572411
c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END,
@@ -2365,8 +2419,7 @@ from
23652419
};
23662420
let expected = indoc!(
23672421
"
2368-
UPDATE t
2369-
SET
2422+
UPDATE t SET
23702423
o = ($5 + $6 + $7 + $8),
23712424
a = CASE WHEN $2 THEN NULL ELSE COALESCE($3, b) END,
23722425
b = CASE WHEN $4 THEN NULL ELSE COALESCE($5, b) END,

‎src/tokenizer.rs‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ pub(crate) enum TokenKind {
8080
ReservedTopLevel,
8181
ReservedTopLevelNoIndent,
8282
ReservedNewline,
83+
ReservedNewlineAfter,
8384
Operator,
8485
OpenParen,
8586
CloseParen,
@@ -418,6 +419,7 @@ fn get_reserved_word_token<'a>(
418419

419420
alt((
420421
get_top_level_reserved_token(last_reserved_top_level_token),
422+
get_newline_after_reserved_token(),
421423
get_newline_reserved_token(last_reserved_token),
422424
get_join_token(),
423425
get_top_level_reserved_token_no_indent,
@@ -524,14 +526,21 @@ fn get_top_level_reserved_token<'a>(
524526
input.to_ascii_uppercase().find(final_word).unwrap_or(0) + final_word.len();
525527
let token = input.next_slice(input_end_pos);
526528

527-
let kind = if token == "EXCEPT"
528-
&& last_reserved_top_level_token.is_some()
529-
&& last_reserved_top_level_token.as_ref().unwrap().value == "SELECT"
530-
{
529+
let kind = match token {
530+
"EXCEPT"
531+
if last_reserved_top_level_token.is_some()
532+
&& last_reserved_top_level_token.as_ref().unwrap().value == "SELECT" =>
531533
// If the query state doesn't allow EXCEPT, treat it as a regular word
532-
TokenKind::Word
533-
} else {
534-
TokenKind::ReservedTopLevel
534+
{
535+
TokenKind::Word
536+
}
537+
"SET"
538+
if last_reserved_top_level_token.is_some()
539+
&& last_reserved_top_level_token.as_ref().unwrap().value == "UPDATE" =>
540+
{
541+
TokenKind::ReservedNewlineAfter
542+
}
543+
_ => TokenKind::ReservedTopLevel,
535544
};
536545

537546
Ok(Token {
@@ -608,6 +617,35 @@ fn get_join_token<'a>() -> impl Parser<&'a str, Token<'a>, ContextError> {
608617
}
609618
}
610619

620+
fn get_newline_after_reserved_token<'a>() -> impl Parser<&'a str, Token<'a>, ContextError> {
621+
move |input: &mut &'a str| {
622+
let uc_input: String = get_uc_words(input, 3);
623+
let mut uc_input = uc_input.as_str();
624+
625+
let mut on_conflict = alt((
626+
terminated("DO NOTHING", end_of_word),
627+
terminated("DO UPDATE SET", end_of_word),
628+
));
629+
630+
let result: Result<&str> = on_conflict.parse_next(&mut uc_input);
631+
632+
if let Ok(token) = result {
633+
let final_word = token.split(' ').next_back().unwrap();
634+
let input_end_pos =
635+
input.to_ascii_uppercase().find(final_word).unwrap() + final_word.len();
636+
let token = input.next_slice(input_end_pos);
637+
let kind = TokenKind::ReservedNewlineAfter;
638+
Ok(Token {
639+
kind,
640+
value: token,
641+
key: None,
642+
})
643+
} else {
644+
Err(ParserError::from_input(input))
645+
}
646+
}
647+
}
648+
611649
fn get_newline_reserved_token<'a>(
612650
last_reserved_token: Option<Token<'a>>,
613651
) -> impl Parser<&'a str, Token<'a>, ContextError> {
@@ -1100,6 +1138,7 @@ fn get_plain_reserved_two_token<'i>(input: &mut &'i str) -> Result<Token<'i>> {
11001138
let result: Result<&str> = alt((
11011139
terminated("CHARACTER SET", end_of_word),
11021140
terminated("ON CONFLICT", end_of_word),
1141+
terminated("ON CONSTRAINT", end_of_word),
11031142
terminated("ON DELETE", end_of_word),
11041143
terminated("ON UPDATE", end_of_word),
11051144
terminated("DISTINCT FROM", end_of_word),

0 commit comments

Comments
 (0)