From 2a8c5c9be906e544e4bf7b5e9f044ffef27ca5c1 Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Wed, 8 Oct 2025 16:11:22 +0200 Subject: [PATCH 1/2] feat!: Support more CREATE/ALTER/DROP statements Fixes #84 --- src/lib.rs | 51 ++++++++++++++++++++++++++++++++++++++++++++---- src/tokenizer.rs | 44 ++++++++++++++++++++++++++++++++--------- 2 files changed, 82 insertions(+), 13 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index a3ad96f..cf98758 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -660,6 +660,46 @@ mod tests { assert_eq!(format(input, &QueryParams::None, &options), expected); } + #[test] + fn it_does_format_drop() { + let input = indoc!( + " + DROP INDEX IF EXISTS idx_a; + DROP INDEX IF EXISTS idx_b; + " + ); + + let options = FormatOptions { + ..Default::default() + }; + + let expected = indoc!( + " + DROP INDEX IF EXISTS + idx_a; + DROP INDEX IF EXISTS + idx_b;" + ); + + assert_eq!(format(input, &QueryParams::None, &options), expected); + + let input = indoc!( + r#" + -- comment + DROP TABLE IF EXISTS "public"."table_name"; + "# + ); + + let expected = indoc!( + r#" + -- comment + DROP TABLE IF EXISTS + "public"."table_name";"# + ); + + assert_eq!(format(input, &QueryParams::None, &options), expected); + } + #[test] fn it_formats_select_query_with_inner_join() { let input = indoc!( @@ -995,8 +1035,12 @@ mod tests { fn it_formats_simple_drop_query() { let input = "DROP TABLE IF EXISTS admin_role;"; let options = FormatOptions::default(); - - assert_eq!(format(input, &QueryParams::None, &options), input); + let output = indoc!( + " + DROP TABLE IF EXISTS + admin_role;" + ); + assert_eq!(format(input, &QueryParams::None, &options), output); } #[test] @@ -1418,8 +1462,7 @@ mod tests { " ALTER TABLE supplier - ALTER COLUMN - supplier_name VARCHAR(100) NOT NULL;" + ALTER COLUMN supplier_name VARCHAR(100) NOT NULL;" ); assert_eq!(format(input, &QueryParams::None, &options), expected); diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 065c20d..a26f26d 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -488,35 +488,61 @@ fn get_top_level_reserved_token<'a>( // First peek at the first character to determine which group to check let first_char = peek(any).parse_next(input)?.to_ascii_uppercase(); + let create_or_replace = ( + "AGGREGATE", + "FUNCTION", + "LANGUAGE", + "PROCEDURE", + "RULE", + "TRIGGER", + "VIEW", + ); + + let alterable_or_droppable = alt((alt(create_or_replace), "TABLE", "INDEX")); + let create_or_replace = alt(create_or_replace); + // Match keywords based on their first letter let result: Result<&str> = match first_char { 'A' => alt(( terminated("ADD", end_of_word), terminated("AFTER", end_of_word), - terminated("ALTER COLUMN", end_of_word), - terminated("ALTER TABLE", end_of_word), + terminated(("ALTER ", alterable_or_droppable).take(), end_of_word), )) .parse_next(&mut uc_input), 'C' => terminated( ( "CREATE ", - opt(alt(( - "UNLOGGED ", + alt(( + create_or_replace, + (opt("UNIQUE "), "INDEX").take(), ( - alt(("GLOBAL ", "LOCAL ")), - opt(alt(("TEMPORARY ", "TEMP "))), + opt(alt(( + "UNLOGGED ", + ( + alt(("GLOBAL ", "LOCAL ")), + opt(alt(("TEMPORARY ", "TEMP "))), + ) + .take(), + ))), + "TABLE", ) .take(), - ))), - "TABLE", + )), ) .take(), end_of_word, ) .parse_next(&mut uc_input), - 'D' => terminated("DELETE FROM", end_of_word).parse_next(&mut uc_input), + 'D' => alt(( + terminated("DELETE FROM", end_of_word), + terminated( + ("DROP ", alterable_or_droppable, opt(" IF EXISTS")).take(), + end_of_word, + ), + )) + .parse_next(&mut uc_input), 'E' => terminated("EXCEPT", end_of_word).parse_next(&mut uc_input), From 8ee53caa8b891c0784afed36b1e5071c302130b3 Mon Sep 17 00:00:00 2001 From: Luca Barbato Date: Tue, 14 Oct 2025 06:04:21 +0200 Subject: [PATCH 2/2] fix: Do not consider ADD a top level token Fixes #105 --- src/lib.rs | 26 ++++++++++++++++++++++++-- src/tokenizer.rs | 1 - 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index cf98758..3d794f0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1468,6 +1468,29 @@ mod tests { assert_eq!(format(input, &QueryParams::None, &options), expected); } + #[test] + fn it_formats_alter_table_add_and_drop() { + let input = r#"ALTER TABLE "public"."event" DROP CONSTRAINT "validate_date", ADD CONSTRAINT "validate_date" CHECK (end_date IS NULL + OR (start_date IS NOT NULL AND end_date > start_date));"#; + + let options = FormatOptions::default(); + let expected = indoc!( + r#" + ALTER TABLE + "public"."event" + DROP CONSTRAINT "validate_date", + ADD CONSTRAINT "validate_date" CHECK ( + end_date IS NULL + OR ( + start_date IS NOT NULL + AND end_date > start_date + ) + );"# + ); + + assert_eq!(format(input, &QueryParams::None, &options), expected); + } + #[test] fn it_recognizes_bracketed_strings() { let inputs = ["[foo JOIN bar]", "[foo ]] JOIN bar]"]; @@ -2239,8 +2262,7 @@ mod tests { -- 自动加载数据到 Hive 分区中 ALTER TABLE sales_data - ADD - PARTITION (sale_year = '2024', sale_month = '08') LOCATION '/user/hive/warehouse/sales_data/2024/08';" + ADD PARTITION (sale_year = '2024', sale_month = '08') LOCATION '/user/hive/warehouse/sales_data/2024/08';" ); assert_eq!(format(input, &QueryParams::None, &options), expected); diff --git a/src/tokenizer.rs b/src/tokenizer.rs index a26f26d..3ddaa6f 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -504,7 +504,6 @@ fn get_top_level_reserved_token<'a>( // Match keywords based on their first letter let result: Result<&str> = match first_char { 'A' => alt(( - terminated("ADD", end_of_word), terminated("AFTER", end_of_word), terminated(("ALTER ", alterable_or_droppable).take(), end_of_word), ))