Recognize UPDATE FOR tokens

Author: lu-zero

src/lib.rs

@@ -456,6 +456,27 @@ mod tests {
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 
+    #[test]
+    fn it_formats_select_with_for_update_of() {
+        let input: &'static str = "SELECT id FROM users WHERE disabled_at IS NULL FOR UPDATE OF users SKIP LOCKED LIMIT 1";
+        let options = FormatOptions::default();
+        let expected = indoc!(
+            "
+            SELECT
+              id
+            FROM
+              users
+            WHERE
+              disabled_at IS NULL
+            FOR UPDATE
+              OF users SKIP LOCKED
+            LIMIT
+              1"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+
     #[test]
     fn it_formats_limit_with_two_comma_separated_values_on_single_line() {
         let input = "LIMIT 5, 10;";

src/tokenizer.rs

@@ -481,7 +481,7 @@ fn get_top_level_reserved_token<'a>(
     last_reserved_top_level_token: Option<Token<'a>>,
 ) -> impl Parser<&'a str, Token<'a>, ContextError> {
     move |input: &mut &'a str| {
-        let uc_input: String = get_uc_words(input, 3);
+        let uc_input: String = get_uc_words(input, 4);
         let mut uc_input = uc_input.as_str();
 
         // First peek at the first character to determine which group to check
@@ -522,6 +522,14 @@ fn get_top_level_reserved_token<'a>(
             'F' => alt((
                 terminated("FETCH FIRST", end_of_word),
                 terminated("FROM", end_of_word),
+                terminated(
+                    (
+                        "FOR ",
+                        alt(("UPDATE", "NO KEY UPDATE", "SHARE", "KEY SHARE")),
+                    )
+                        .take(),
+                    end_of_word,
+                ),
             ))
             .parse_next(&mut uc_input),