diff --git a/fuzz/fuzz_seeds/postgres_2778_row_locking.sql b/fuzz/fuzz_seeds/postgres_2778_row_locking.sql new file mode 100644 index 0000000000..4d712eebe3 --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_2778_row_locking.sql @@ -0,0 +1 @@ +SELECT id FROM t WHERE id = 1 FOR NO KEY UPDATE diff --git a/src/ast/query.rs b/src/ast/query.rs index e2364dbe9a..40cea358ad 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -3658,6 +3658,10 @@ pub enum LockType { Share, /// `UPDATE` lock (exclusive/update lock). Update, + /// `NO KEY UPDATE` lock. + NoKeyUpdate, + /// `KEY SHARE` lock. + KeyShare, } impl fmt::Display for LockType { @@ -3665,6 +3669,8 @@ impl fmt::Display for LockType { let select_lock = match self { LockType::Share => "SHARE", LockType::Update => "UPDATE", + LockType::NoKeyUpdate => "NO KEY UPDATE", + LockType::KeyShare => "KEY SHARE", }; write!(f, "{select_lock}") } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 11accc3996..3a26a99307 100644 --- a/src/parser/mod.rs +++ b/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 { - 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)?) diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index be2c7e9563..2c064ee738 100644 --- a/tests/sqlparser_postgres.rs +++ b/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}" + ); + } +}