diff --git a/src/formatter.rs b/src/formatter.rs index 781d06d..34323cc 100644 --- a/src/formatter.rs +++ b/src/formatter.rs @@ -233,7 +233,7 @@ impl<'a> Formatter<'a> { self.add_new_line(query); } query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value))); - if newline_after { + if newline_after && token.alias != "CREATE" { self.indentation.increase_top_level(span_info); self.add_new_line(query); } else { diff --git a/src/lib.rs b/src/lib.rs index 0f874e1..f17c8d4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -838,6 +838,25 @@ mod tests { assert_eq!(format(input, &QueryParams::None, &options), expected); } + #[test] + fn it_formats_full_delete_query() { + let input = + "DELETE FROM Customers USING Phonebook WHERE CustomerName='Alfred' AND Phone=5002132;"; + let options = FormatOptions::default(); + let expected = indoc!( + " + DELETE FROM + Customers + USING + Phonebook + WHERE + CustomerName = 'Alfred' + AND Phone = 5002132;" + ); + + assert_eq!(format(input, &QueryParams::None, &options), expected); + } + #[test] fn it_formats_simple_drop_query() { let input = "DROP TABLE IF EXISTS admin_role;"; diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 78ca506..ab813f6 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -69,6 +69,8 @@ pub(crate) struct Token<'a> { pub value: &'a str, // Only used for placeholder--there is a reason this isn't on the enum pub key: Option>, + /// Used to group the behaviour of variants of tokens + pub alias: &'a str, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -166,6 +168,7 @@ fn get_type_specifier_token<'i>( kind: TokenKind::TypeSpecifier, value: token, key: None, + alias: token, }) } } @@ -176,6 +179,7 @@ fn get_whitespace_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::Whitespace, value: token, key: None, + alias: token, }) } @@ -192,6 +196,7 @@ fn get_comment_token<'i>(input: &mut &'i str) -> Result> { kind, value: token, key: None, + alias: token, }) } @@ -250,6 +255,7 @@ fn get_string_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::String, value: token, key: None, + alias: token, }) } @@ -269,6 +275,7 @@ fn get_placeholder_string_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::String, value: token, key: None, + alias: token, }) } @@ -279,6 +286,7 @@ fn get_open_paren_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::OpenParen, value: token, key: None, + alias: token, }) } @@ -289,6 +297,7 @@ fn get_close_paren_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::CloseParen, value: token, key: None, + alias: token, }) } @@ -332,6 +341,7 @@ fn get_indexed_placeholder_token<'i>(input: &mut &'i str) -> Result> { } else { None }, + alias: token, }) } @@ -350,6 +360,7 @@ fn get_ident_named_placeholder_token<'i>(input: &mut &'i str) -> Result(input: &mut &'i str) -> Result(input: &mut &'i str) -> Result> { kind: TokenKind::Number, value: token, key: None, + alias: token, }) } @@ -464,6 +477,24 @@ fn get_top_level_reserved_token<'a>( )) .parse_next(&mut uc_input), + 'C' => terminated( + ( + "CREATE ", + opt(alt(( + "UNLOGGED ", + ( + alt(("GLOBAL ", "LOCAL ")), + opt(alt(("TEMPORARY ", "TEMP "))), + ) + .take(), + ))), + "TABLE", + ) + .take(), + end_of_word, + ) + .parse_next(&mut uc_input), + 'D' => terminated("DELETE FROM", end_of_word).parse_next(&mut uc_input), 'E' => terminated("EXCEPT", end_of_word).parse_next(&mut uc_input), @@ -490,7 +521,11 @@ fn get_top_level_reserved_token<'a>( 'L' => terminated("LIMIT", end_of_word).parse_next(&mut uc_input), - 'M' => terminated("MODIFY", end_of_word).parse_next(&mut uc_input), + 'M' => alt(( + terminated("MODIFY", end_of_word), + terminated("MERGE INTO", end_of_word), + )) + .parse_next(&mut uc_input), 'O' => alt(( terminated("ORDER BY", end_of_word), @@ -512,7 +547,11 @@ fn get_top_level_reserved_token<'a>( )) .parse_next(&mut uc_input), - 'U' => terminated("UPDATE", end_of_word).parse_next(&mut uc_input), + 'U' => alt(( + terminated("UPDATE", end_of_word), + terminated("USING", end_of_word), + )) + .parse_next(&mut uc_input), 'V' => terminated("VALUES", end_of_word).parse_next(&mut uc_input), @@ -529,27 +568,35 @@ fn get_top_level_reserved_token<'a>( if let Ok(token) = result { let token = finalize(input, token); - let kind = match token { - "EXCEPT" - if last_reserved_top_level_token.is_some() - && last_reserved_top_level_token.as_ref().unwrap().value == "SELECT" => - // If the query state doesn't allow EXCEPT, treat it as a regular word + let kind = match ( + token, + last_reserved_top_level_token.as_ref().map(|v| v.alias), + ) { + ("EXCEPT", Some("SELECT")) => + // If the query state doesn't allow EXCEPT, treat it as a reserved word { - TokenKind::Word + TokenKind::Reserved } - "SET" - if last_reserved_top_level_token.is_some() - && last_reserved_top_level_token.as_ref().unwrap().value == "UPDATE" => - { - TokenKind::ReservedNewlineAfter + ("SET", Some("UPDATE")) => TokenKind::ReservedNewlineAfter, + ("USING", v) if v != Some("MERGE INTO") && v != Some("DELETE FROM") => { + TokenKind::Reserved } _ => TokenKind::ReservedTopLevel, }; + let alias = if token.starts_with("CREATE") { + "CREATE" + } else if token.starts_with("SELECT") { + "SELECT" + } else { + token + }; + Ok(Token { kind, value: token, key: None, + alias, }) } else { Err(ParserError::from_input(input)) @@ -613,6 +660,7 @@ fn get_join_token<'a>() -> impl Parser<&'a str, Token<'a>, ContextError> { kind, value: token, key: None, + alias: token, }) } else { Err(ParserError::from_input(input)) @@ -633,10 +681,12 @@ fn get_newline_after_reserved_token<'a>() -> impl Parser<&'a str, Token<'a>, Con let result: Result<&str> = on_conflict.parse_next(&mut uc_input); if let Ok(token) = result { + let value = finalize(input, token); Ok(Token { kind: TokenKind::ReservedNewlineAfter, - value: finalize(input, token), + value, key: None, + alias: value, }) } else { Err(ParserError::from_input(input)) @@ -692,6 +742,7 @@ fn get_newline_reserved_token<'a>( kind, value: token, key: None, + alias: token, }) } else { Err(ParserError::from_input(input)) @@ -716,10 +767,12 @@ fn get_top_level_reserved_token_no_indent<'i>(input: &mut &'i str) -> Result(input: &mut &'i str) -> Result> { terminated("TRUNCATE", end_of_word), terminated("TYPE", end_of_word), terminated("TYPES", end_of_word), + terminated("TBLPROPERTIES", end_of_word), )) .parse_next(&mut uc_input), @@ -1090,7 +1144,6 @@ fn get_plain_reserved_one_token<'i>(input: &mut &'i str) -> Result> { terminated("UNSIGNED", end_of_word), terminated("USAGE", end_of_word), terminated("USE", end_of_word), - terminated("USING", end_of_word), )) .parse_next(&mut uc_input), @@ -1118,6 +1171,7 @@ fn get_plain_reserved_one_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::Reserved, value: token, key: None, + alias: token, }) } else { Err(ParserError::from_input(input)) @@ -1134,13 +1188,16 @@ fn get_plain_reserved_two_token<'i>(input: &mut &'i str) -> Result> { terminated("ON DELETE", end_of_word), terminated("ON UPDATE", end_of_word), terminated("DISTINCT FROM", end_of_word), + terminated("PARTITIONED BY", end_of_word), )) .parse_next(&mut uc_input); if let Ok(token) = result { + let value = finalize(input, token); Ok(Token { kind: TokenKind::Reserved, - value: finalize(input, token), + value, key: None, + alias: value, }) } else { Err(ParserError::from_input(input)) @@ -1154,6 +1211,7 @@ fn get_word_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::Word, value: token, key: None, + alias: token, }) } @@ -1168,6 +1226,7 @@ fn get_operator_token<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::Operator, value: token, key: None, + alias: token, }) .parse_next(input) } @@ -1179,6 +1238,7 @@ fn get_any_other_char<'i>(input: &mut &'i str) -> Result> { kind: TokenKind::Operator, value: token, key: None, + alias: token, }) }