PostgreSQL: Support FOR NO KEY UPDATE and FOR KEY SHARE

fuzz/fuzz_seeds/postgres_2778_row_locking.sql

@@ -0,0 +1 @@
+SELECT id FROM t WHERE id = 1 FOR NO KEY UPDATE

src/ast/query.rs

@@ -3658,13 +3658,19 @@ pub enum LockType {
     Share,
     /// `UPDATE` lock (exclusive/update lock).
     Update,
+    /// `NO KEY UPDATE` lock.
+    NoKeyUpdate,
+    /// `KEY SHARE` lock.
+    KeyShare,
 }
 
 impl fmt::Display for LockType {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let select_lock = match self {
             LockType::Share => "SHARE",
             LockType::Update => "UPDATE",
+            LockType::NoKeyUpdate => "NO KEY UPDATE",
+            LockType::KeyShare => "KEY SHARE",
         };
         write!(f, "{select_lock}")
     }

src/parser/mod.rs

@@ -20713,14 +20713,22 @@ impl<'a> Parser<'a> {
         })
     }
 
-    /// Parse a FOR UPDATE/FOR SHARE clause
+    /// Parse a `FOR` row-locking clause.
     pub fn parse_lock(&mut self) -> Result<LockClause, ParserError> {
-        let lock_type = match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
-            Keyword::UPDATE => LockType::Update,
-            Keyword::SHARE => LockType::Share,
-            unexpected_keyword => return Err(ParserError::ParserError(
-                format!("Internal parser error: expected any of {{UPDATE, SHARE}}, got {unexpected_keyword:?}"),
-            )),
+        let lock_type = if self.parse_keyword(Keyword::NO) {
+            self.expect_keywords(&[Keyword::KEY, Keyword::UPDATE])?;
+            LockType::NoKeyUpdate
+        } else if self.parse_keyword(Keyword::KEY) {
+            self.expect_keyword(Keyword::SHARE)?;
+            LockType::KeyShare
+        } else {
+            match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
+                Keyword::UPDATE => LockType::Update,
+                Keyword::SHARE => LockType::Share,
+                unexpected_keyword => return Err(ParserError::ParserError(
+                    format!("Internal parser error: expected any of {{UPDATE, SHARE}}, got {unexpected_keyword:?}"),
+                )),
+            }
         };
         let of = if self.parse_keyword(Keyword::OF) {
             Some(self.parse_object_name(false)?)

tests/sqlparser_postgres.rs

@@ -12291,3 +12291,77 @@ fn parse_create_table_partition_of_untyped_columns() {
         );
     }
 }
+
+#[test]
+fn parse_for_no_key_update_and_key_share() {
+    for (sql, expected_lock_type, expected_of, expected_nonblock) in [
+        (
+            "SELECT id FROM t WHERE id = 1 FOR NO KEY UPDATE",
+            LockType::NoKeyUpdate,
+            None,
+            None,
+        ),
+        (
+            "SELECT id FROM t WHERE id = 1 FOR NO KEY UPDATE OF t SKIP LOCKED",
+            LockType::NoKeyUpdate,
+            Some("t"),
+            Some(NonBlock::SkipLocked),
+        ),
+        (
+            "SELECT id FROM t WHERE id = 1 FOR KEY SHARE",
+            LockType::KeyShare,
+            None,
+            None,
+        ),
+        (
+            "SELECT id FROM t WHERE id = 1 FOR KEY SHARE OF t NOWAIT",
+            LockType::KeyShare,
+            Some("t"),
+            Some(NonBlock::Nowait),
+        ),
+    ] {
+        let statement = pg_and_generic().verified_stmt(sql);
+        assert_eq!(statement.to_string(), sql);
+
+        let Statement::Query(query) = statement else {
+            panic!("expected a query for {sql}");
+        };
+        assert_eq!(query.locks.len(), 1, "{sql}");
+        let lock = &query.locks[0];
+        assert_eq!(lock.lock_type, expected_lock_type, "{sql}");
+        assert_eq!(
+            lock.of.as_ref().map(ToString::to_string).as_deref(),
+            expected_of,
+            "{sql}"
+        );
+        assert_eq!(lock.nonblock, expected_nonblock, "{sql}");
+    }
+
+    assert_ne!(LockType::NoKeyUpdate, LockType::Update);
+    assert_ne!(LockType::KeyShare, LockType::Share);
+}
+
+#[test]
+fn reject_incomplete_for_lock_strength() {
+    for (sql, expected) in [
+        ("SELECT * FROM t FOR NO KEY", "Expected: UPDATE, found: EOF"),
+        (
+            "SELECT * FROM t FOR NO KEY SHARE",
+            "Expected: UPDATE, found: SHARE",
+        ),
+        ("SELECT * FROM t FOR KEY", "Expected: SHARE, found: EOF"),
+        (
+            "SELECT * FROM t FOR KEY UPDATE",
+            "Expected: SHARE, found: UPDATE",
+        ),
+    ] {
+        assert_eq!(
+            pg_and_generic()
+                .parse_sql_statements(sql)
+                .unwrap_err()
+                .to_string(),
+            format!("sql parser error: {expected}"),
+            "{sql}"
+        );
+    }
+}