From 6caf0b8d365cef594765901d5220445ccc655425 Mon Sep 17 00:00:00 2001
From: Will Mortensen <will@extrahop.com>
Date: Mon, 16 Jan 2023 20:51:26 -0800
Subject: [PATCH] Add WAIT ONLY option to LOCK

Rather than actually taking any locks on the table(s), it simply waits
for conflicting lockers using the existing WaitForLockersMultiple()
function in the lock manager (previously used only by concurrent index
operations). As when actually taking the lock, it doesn't wait for any
conflicting locks acquired after it initially determines the set of
conflicting transactions.

Currently it's not supported with views, since they would require more
locking to gather the locktags.

The syntax allows combining it with NO WAIT, which would perhaps be
useful to simply check for conflicts, but this is not yet implemented.
(NOWAIT + immediately releasing the lock already accomplishes roughly
the same thing.)

Unlike other forms of LOCK, WAIT ONLY is allowed outside a transaction
block, since it makes perfect sense to wait and then e.g. SELECT new
data.

Regardless of the specified locking mode, only SELECT permissions are
required on the table(s).

XXX: docs have not been updated yet.
---
 src/backend/commands/lockcmds.c               | 83 +++++++++++++++++--
 src/backend/parser/gram.y                     | 14 +++-
 src/backend/tcop/utility.c                    | 24 ++++--
 src/include/nodes/parsenodes.h                |  1 +
 src/include/parser/kwlist.h                   |  1 +
 .../isolation/expected/deadlock-wait-only.out | 12 +++
 src/test/isolation/expected/wait-only.out     | 78 +++++++++++++++++
 src/test/isolation/isolation_schedule         |  2 +
 .../isolation/specs/deadlock-wait-only.spec   | 23 +++++
 src/test/isolation/specs/wait-only.spec       | 47 +++++++++++
 src/test/regress/expected/lock.out            | 30 +++++++
 src/test/regress/sql/lock.sql                 | 32 +++++++
 12 files changed, 329 insertions(+), 18 deletions(-)
 create mode 100644 src/test/isolation/expected/deadlock-wait-only.out
 create mode 100644 src/test/isolation/expected/wait-only.out
 create mode 100644 src/test/isolation/specs/deadlock-wait-only.spec
 create mode 100644 src/test/isolation/specs/wait-only.spec

diff --git a/src/backend/commands/lockcmds.c b/src/backend/commands/lockcmds.c
index 43c7d7f4bb..90976abdb1 100644
--- a/src/backend/commands/lockcmds.c
+++ b/src/backend/commands/lockcmds.c
@@ -16,6 +16,7 @@
 
 #include "access/table.h"
 #include "access/xact.h"
+#include "catalog/catalog.h"
 #include "catalog/namespace.h"
 #include "catalog/pg_inherits.h"
 #include "commands/lockcmds.h"
@@ -29,7 +30,8 @@
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
 
-static void LockTableRecurse(Oid reloid, LOCKMODE lockmode, bool nowait);
+static void LockTableRecurse(Oid reloid, LOCKMODE lockmode, bool nowait,
+							 List **locktags_p);
 static AclResult LockTableAclCheck(Oid reloid, LOCKMODE lockmode, Oid userid);
 static void RangeVarCallbackForLockTable(const RangeVar *rv, Oid relid,
 										 Oid oldrelid, void *arg);
@@ -43,6 +45,34 @@ void
 LockTableCommand(LockStmt *lockstmt)
 {
 	ListCell   *p;
+	LOCKMODE	lockmode;
+	LOCKMODE   	waitmode;
+	List	   *waitlocktags = NIL;
+	List	  **waitlocktags_p;
+
+	if (lockstmt->waitonly && lockstmt->nowait)
+		/*
+		 * this could be defined to check and error if there are conflicting
+		 * lockers, but it seems unclear if that would be useful, since
+		 * LOCK ... NOWAIT + immediate unlock would do nearly the same thing
+		 */
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				errmsg("NOWAIT is not supported with WAIT ONLY")));
+
+
+	if (lockstmt->waitonly)
+	{
+		lockmode = NoLock;
+		waitmode = lockstmt->mode;
+		waitlocktags_p = &waitlocktags;
+	}
+	else
+	{
+		lockmode = lockstmt->mode;
+		waitmode = NoLock;
+		waitlocktags_p = NULL;
+	}
 
 	/*
 	 * Iterate over the list and process the named relations one at a time
@@ -53,16 +83,37 @@ LockTableCommand(LockStmt *lockstmt)
 		bool		recurse = rv->inh;
 		Oid			reloid;
 
-		reloid = RangeVarGetRelidExtended(rv, lockstmt->mode,
+		reloid = RangeVarGetRelidExtended(rv, lockmode,
 										  lockstmt->nowait ? RVR_NOWAIT : 0,
 										  RangeVarCallbackForLockTable,
-										  (void *) &lockstmt->mode);
+										  (void *) &lockmode);
+		if (waitmode != NoLock)
+		{
+			Oid			dbid;
+			LOCKTAG	   *heaplocktag = palloc_object(LOCKTAG);
+
+			if (IsSharedRelation(reloid))
+				dbid = InvalidOid;
+			else
+				dbid = MyDatabaseId;
+			SET_LOCKTAG_RELATION(*heaplocktag, dbid, reloid);
+			waitlocktags = lappend(waitlocktags, heaplocktag);
+		}
 
 		if (get_rel_relkind(reloid) == RELKIND_VIEW)
-			LockViewRecurse(reloid, lockstmt->mode, lockstmt->nowait, NIL);
+		{
+			if (lockstmt->waitonly || lockmode == NoLock)
+				ereport(ERROR,
+						(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+						errmsg("WAIT ONLY is not supported with views")));
+			LockViewRecurse(reloid, lockmode, lockstmt->nowait, NIL);
+		}
 		else if (recurse)
-			LockTableRecurse(reloid, lockstmt->mode, lockstmt->nowait);
+			LockTableRecurse(reloid, lockmode, lockstmt->nowait,
+							 waitlocktags_p);
 	}
+	if (waitmode != NoLock)
+		WaitForLockersMultiple(waitlocktags, waitmode, false);
 }
 
 /*
@@ -116,7 +167,7 @@ RangeVarCallbackForLockTable(const RangeVar *rv, Oid relid, Oid oldrelid,
  * parent which is enough.
  */
 static void
-LockTableRecurse(Oid reloid, LOCKMODE lockmode, bool nowait)
+LockTableRecurse(Oid reloid, LOCKMODE lockmode, bool nowait, List **locktags_p)
 {
 	List	   *children;
 	ListCell   *lc;
@@ -126,11 +177,26 @@ LockTableRecurse(Oid reloid, LOCKMODE lockmode, bool nowait)
 	foreach(lc, children)
 	{
 		Oid			childreloid = lfirst_oid(lc);
+		Oid			dbid;
+		LOCKTAG	   *heaplocktag;
 
-		/* Parent already locked. */
+		/* Parent already handled. */
 		if (childreloid == reloid)
 			continue;
 
+		if (locktags_p != NULL)
+		{
+			heaplocktag = palloc_object(LOCKTAG);
+			if (IsSharedRelation(childreloid))
+				dbid = InvalidOid;
+			else
+				dbid = MyDatabaseId;
+			SET_LOCKTAG_RELATION(*heaplocktag, dbid, childreloid);
+			*locktags_p = lappend(*locktags_p, heaplocktag);
+		}
+
+		if (lockmode == NoLock)
+			continue;
 		if (!nowait)
 			LockRelationOid(childreloid, lockmode);
 		else if (!ConditionalLockRelationOid(childreloid, lockmode))
@@ -229,7 +295,8 @@ LockViewRecurse_walker(Node *node, LockViewRecurse_context *context)
 				LockViewRecurse(relid, context->lockmode, context->nowait,
 								context->ancestor_views);
 			else if (rte->inh)
-				LockTableRecurse(relid, context->lockmode, context->nowait);
+				LockTableRecurse(relid, context->lockmode, context->nowait,
+								 NULL);
 		}
 
 		return query_tree_walker(query,
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index a0138382a1..db07eadc68 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -351,7 +351,7 @@ static Node *makeRecursiveViewSelect(char *relname, List *aliases, Node *query);
 %type <defelt>	drop_option
 %type <boolean>	opt_or_replace opt_no
 				opt_grant_grant_option
-				opt_nowait opt_if_exists opt_with_data
+				opt_nowait opt_waitonly opt_if_exists opt_with_data
 				opt_transaction_chain
 %type <list>	grant_role_opt_list
 %type <defelt>	grant_role_opt
@@ -755,7 +755,7 @@ static Node *makeRecursiveViewSelect(char *relname, List *aliases, Node *query);
 	VACUUM VALID VALIDATE VALIDATOR VALUE_P VALUES VARCHAR VARIADIC VARYING
 	VERBOSE VERSION_P VIEW VIEWS VOLATILE
 
-	WHEN WHERE WHITESPACE_P WINDOW WITH WITHIN WITHOUT WORK WRAPPER WRITE
+	WAIT WHEN WHERE WHITESPACE_P WINDOW WITH WITHIN WITHOUT WORK WRAPPER WRITE
 
 	XML_P XMLATTRIBUTES XMLCONCAT XMLELEMENT XMLEXISTS XMLFOREST XMLNAMESPACES
 	XMLPARSE XMLPI XMLROOT XMLSERIALIZE XMLTABLE
@@ -12117,13 +12117,14 @@ using_clause:
  *
  *****************************************************************************/
 
-LockStmt:	LOCK_P opt_table relation_expr_list opt_lock opt_nowait
+LockStmt:	LOCK_P opt_table relation_expr_list opt_lock opt_nowait opt_waitonly
 				{
 					LockStmt   *n = makeNode(LockStmt);
 
 					n->relations = $3;
 					n->mode = $4;
 					n->nowait = $5;
+					n->waitonly = $6;
 					$$ = (Node *) n;
 				}
 		;
@@ -12146,6 +12147,11 @@ opt_nowait:	NOWAIT							{ $$ = true; }
 			| /*EMPTY*/						{ $$ = false; }
 		;
 
+opt_waitonly:
+			WAIT ONLY						{ $$ = true; }
+			| /*EMPTY*/						{ $$ = false; }
+		;
+
 opt_nowait_or_skip:
 			NOWAIT							{ $$ = LockWaitError; }
 			| SKIP LOCKED					{ $$ = LockWaitSkip; }
@@ -17010,6 +17016,7 @@ unreserved_keyword:
 			| VIEW
 			| VIEWS
 			| VOLATILE
+			| WAIT
 			| WHITESPACE_P
 			| WITHIN
 			| WITHOUT
@@ -17622,6 +17629,7 @@ bare_label_keyword:
 			| VIEW
 			| VIEWS
 			| VOLATILE
+			| WAIT
 			| WHEN
 			| WHITESPACE_P
 			| WORK
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
index c7d9d96b45..49f0a99943 100644
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -354,7 +354,7 @@ ClassifyUtilityCommandAsReadOnly(Node *parsetree)
 				 * restrictions here must match those in
 				 * LockAcquireExtended().
 				 */
-				if (stmt->mode > RowExclusiveLock)
+				if (!stmt->waitonly && stmt->mode > RowExclusiveLock)
 					return COMMAND_OK_IN_READ_ONLY_TXN;
 				else
 					return COMMAND_IS_STRICTLY_READ_ONLY;
@@ -932,13 +932,23 @@ standard_ProcessUtility(PlannedStmt *pstmt,
 			break;
 
 		case T_LockStmt:
+			{
+				LockStmt *stmt = (LockStmt *) parsetree;
 
-			/*
-			 * Since the lock would just get dropped immediately, LOCK TABLE
-			 * outside a transaction block is presumed to be user error.
-			 */
-			RequireTransactionBlock(isTopLevel, "LOCK TABLE");
-			LockTableCommand((LockStmt *) parsetree);
+				if (!stmt->waitonly)
+				{
+					/*
+					 * Since the lock would just get dropped immediately, and
+					 * simply waiting is better done with WAIT ONLY, LOCK TABLE
+					 * without WAIT ONLY outside a transaction block is presumed
+					 * to be user error.
+					 *
+					 * XXX: the error should clarify that WAIT ONLY is allowed?
+					 */
+					RequireTransactionBlock(isTopLevel, "LOCK TABLE");
+				}
+				LockTableCommand(stmt);
+			}
 			break;
 
 		case T_ConstraintsSetStmt:
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 89335d95e7..a44eaf3aa3 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -3631,6 +3631,7 @@ typedef struct LockStmt
 	List	   *relations;		/* relations to lock */
 	int			mode;			/* lock mode */
 	bool		nowait;			/* no wait mode */
+	bool		waitonly;		/* wait only mode */
 } LockStmt;
 
 /* ----------------------
diff --git a/src/include/parser/kwlist.h b/src/include/parser/kwlist.h
index bb36213e6f..1cd1ab6dfd 100644
--- a/src/include/parser/kwlist.h
+++ b/src/include/parser/kwlist.h
@@ -461,6 +461,7 @@ PG_KEYWORD("version", VERSION_P, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("view", VIEW, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("views", VIEWS, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("volatile", VOLATILE, UNRESERVED_KEYWORD, BARE_LABEL)
+PG_KEYWORD("wait", WAIT, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("when", WHEN, RESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("where", WHERE, RESERVED_KEYWORD, AS_LABEL)
 PG_KEYWORD("whitespace", WHITESPACE_P, UNRESERVED_KEYWORD, BARE_LABEL)
diff --git a/src/test/isolation/expected/deadlock-wait-only.out b/src/test/isolation/expected/deadlock-wait-only.out
new file mode 100644
index 0000000000..78b4962fa1
--- /dev/null
+++ b/src/test/isolation/expected/deadlock-wait-only.out
@@ -0,0 +1,12 @@
+Parsed test spec with 2 sessions
+
+starting permutation: s1re s2as s2swo s1aewo s1c s2c
+step s1re: LOCK TABLE a1 IN ROW EXCLUSIVE MODE;
+step s2as: LOCK TABLE a1 IN ACCESS SHARE MODE;
+step s2swo: LOCK TABLE a1 IN SHARE MODE WAIT ONLY; <waiting ...>
+step s1aewo: LOCK TABLE a1 IN ACCESS EXCLUSIVE MODE WAIT ONLY; <waiting ...>
+step s1aewo: <... completed>
+step s2swo: <... completed>
+ERROR:  deadlock detected
+step s1c: COMMIT;
+step s2c: COMMIT;
diff --git a/src/test/isolation/expected/wait-only.out b/src/test/isolation/expected/wait-only.out
new file mode 100644
index 0000000000..2358ef653b
--- /dev/null
+++ b/src/test/isolation/expected/wait-only.out
@@ -0,0 +1,78 @@
+Parsed test spec with 3 sessions
+
+starting permutation: w1in rlwo w1c rsel rc w2c
+step w1in: INSERT INTO a1 VALUES (DEFAULT);
+step rlwo: LOCK TABLE a1 IN SHARE MODE WAIT ONLY; <waiting ...>
+step w1c: COMMIT;
+step rlwo: <... completed>
+step rsel: SELECT id from a1;
+id
+--
+ 1
+(1 row)
+
+step rc: COMMIT;
+step w2c: COMMIT;
+
+starting permutation: w1in w1c rlwo rsel rc w2c
+step w1in: INSERT INTO a1 VALUES (DEFAULT);
+step w1c: COMMIT;
+step rlwo: LOCK TABLE a1 IN SHARE MODE WAIT ONLY;
+step rsel: SELECT id from a1;
+id
+--
+ 1
+(1 row)
+
+step rc: COMMIT;
+step w2c: COMMIT;
+
+starting permutation: w1in rlwo w2in w2c w1c rsel rc
+step w1in: INSERT INTO a1 VALUES (DEFAULT);
+step rlwo: LOCK TABLE a1 IN SHARE MODE WAIT ONLY; <waiting ...>
+step w2in: INSERT INTO a1 VALUES (DEFAULT);
+step w2c: COMMIT;
+step w1c: COMMIT;
+step rlwo: <... completed>
+step rsel: SELECT id from a1;
+id
+--
+ 1
+ 2
+(2 rows)
+
+step rc: COMMIT;
+
+starting permutation: w1in rsv rl w2in w1c rrb w2c rsel rc
+step w1in: INSERT INTO a1 VALUES (DEFAULT);
+step rsv: SAVEPOINT foo;
+step rl: LOCK TABLE a1 IN SHARE MODE; <waiting ...>
+step w2in: INSERT INTO a1 VALUES (DEFAULT); <waiting ...>
+step w1c: COMMIT;
+step rl: <... completed>
+step rrb: ROLLBACK TO foo;
+step w2in: <... completed>
+step w2c: COMMIT;
+step rsel: SELECT id from a1;
+id
+--
+ 1
+ 2
+(2 rows)
+
+step rc: COMMIT;
+
+starting permutation: w1in rlwo w2in w1c rsel rc w2c
+step w1in: INSERT INTO a1 VALUES (DEFAULT);
+step rlwo: LOCK TABLE a1 IN SHARE MODE WAIT ONLY; <waiting ...>
+step w2in: INSERT INTO a1 VALUES (DEFAULT);
+step w1c: COMMIT;
+step rlwo: <... completed>
+step rsel: SELECT id from a1;
+id
+--
+ 1
+(1 row)
+
+step rc: COMMIT;
+step w2c: COMMIT;
diff --git a/src/test/isolation/isolation_schedule b/src/test/isolation/isolation_schedule
index c11dc9a420..b83e371aec 100644
--- a/src/test/isolation/isolation_schedule
+++ b/src/test/isolation/isolation_schedule
@@ -5,6 +5,7 @@ test: read-write-unique
 test: read-write-unique-2
 test: read-write-unique-3
 test: read-write-unique-4
+test: wait-only
 test: simple-write-skew
 test: receipt-report
 test: temporal-range-integrity
@@ -20,6 +21,7 @@ test: index-only-scan
 test: predicate-lock-hot-tuple
 test: update-conflict-out
 test: deadlock-simple
+test: deadlock-wait-only
 test: deadlock-hard
 test: deadlock-soft
 test: deadlock-soft-2
diff --git a/src/test/isolation/specs/deadlock-wait-only.spec b/src/test/isolation/specs/deadlock-wait-only.spec
new file mode 100644
index 0000000000..0efca38d60
--- /dev/null
+++ b/src/test/isolation/specs/deadlock-wait-only.spec
@@ -0,0 +1,23 @@
+setup
+{
+  CREATE TABLE a1 ();
+}
+
+teardown
+{
+  DROP TABLE a1;
+}
+
+session s1
+setup		{ BEGIN; }
+step s1re	  { LOCK TABLE a1 IN ROW EXCLUSIVE MODE; }
+step s1aewo	{ LOCK TABLE a1 IN ACCESS EXCLUSIVE MODE WAIT ONLY; }
+step s1c	  { COMMIT; }
+
+session s2
+setup		{ BEGIN; }
+step s2as		{ LOCK TABLE a1 IN ACCESS SHARE MODE; }
+step s2swo	{ LOCK TABLE a1 IN SHARE MODE WAIT ONLY; }
+step s2c		{ COMMIT; }
+
+permutation s1re s2as s2swo s1aewo s1c s2c
diff --git a/src/test/isolation/specs/wait-only.spec b/src/test/isolation/specs/wait-only.spec
new file mode 100644
index 0000000000..f91ac3ba18
--- /dev/null
+++ b/src/test/isolation/specs/wait-only.spec
@@ -0,0 +1,47 @@
+setup
+{
+  CREATE TABLE a1 (id bigserial);
+}
+
+teardown
+{
+  DROP TABLE a1;
+}
+
+# use READ COMMITTED so we can observe the effects of a committed INSERT after
+# waiting
+
+session writer1
+setup		{ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED; }
+step w1in	{ INSERT INTO a1 VALUES (DEFAULT); }
+step w1c	{ COMMIT; }
+
+session writer2
+setup		{ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED; }
+step w2in	{ INSERT INTO a1 VALUES (DEFAULT); }
+step w2c	{ COMMIT; }
+
+session reader
+setup		{ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED; }
+step rsv	{ SAVEPOINT foo; }
+step rl   { LOCK TABLE a1 IN SHARE MODE; }
+step rrb  { ROLLBACK TO foo; }
+step rlwo	{ LOCK TABLE a1 IN SHARE MODE WAIT ONLY; }
+step rsel	{ SELECT id from a1; }
+step rc		{ COMMIT; }
+
+# reader waits for writer1 (writer2 no-op)
+permutation w1in rlwo w1c rsel rc w2c
+
+# no waiting if writer1 committed (writer2 no-op)
+permutation w1in w1c rlwo rsel rc w2c
+
+# reader waiting for writer1 doesn't block writer2...
+permutation w1in rlwo w2in w2c w1c rsel rc
+# ...while actually taking the lock does block writer2 (even if we release it
+# ASAP)
+permutation w1in rsv rl w2in w1c rrb w2c rsel rc
+
+# reader doesn't wait for lock newly acquired by writer2 while waiting for
+# writer1
+permutation w1in rlwo w2in w1c rsel rc w2c
diff --git a/src/test/regress/expected/lock.out b/src/test/regress/expected/lock.out
index ad137d3645..98de54d68b 100644
--- a/src/test/regress/expected/lock.out
+++ b/src/test/regress/expected/lock.out
@@ -41,6 +41,22 @@ LOCK TABLE lock_tbl1 IN SHARE ROW EXCLUSIVE MODE NOWAIT;
 LOCK TABLE lock_tbl1 IN EXCLUSIVE MODE NOWAIT;
 LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE NOWAIT;
 ROLLBACK;
+-- Try using WAIT ONLY along with valid options.
+BEGIN TRANSACTION;
+LOCK TABLE lock_tbl1 IN ACCESS SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ROW SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ROW EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE UPDATE EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE ROW EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+ROLLBACK;
+-- WAIT ONLY is allowed outside a transaction
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+-- NOWAIT + WAIT ONLY is not supported (yet?)
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE NOWAIT WAIT ONLY;
+ERROR:  NOWAIT is not supported with WAIT ONLY
 -- Verify that we can lock views.
 BEGIN TRANSACTION;
 LOCK TABLE lock_view1 IN EXCLUSIVE MODE;
@@ -138,8 +154,22 @@ ROLLBACK;
 CREATE TABLE lock_tbl2 (b BIGINT) INHERITS (lock_tbl1);
 CREATE TABLE lock_tbl3 () INHERITS (lock_tbl2);
 BEGIN TRANSACTION;
+LOCK TABLE lock_tbl1 * IN ACCESS EXCLUSIVE MODE WAIT ONLY;
 LOCK TABLE lock_tbl1 * IN ACCESS EXCLUSIVE MODE;
 ROLLBACK;
+-- WAIT ONLY requires SELECT permissions regardless of lock mode
+-- fail without permissions
+SET ROLE regress_rol_lock1;
+BEGIN;
+LOCK TABLE ONLY lock_tbl1 IN ACCESS SHARE MODE WAIT ONLY;
+ERROR:  permission denied for table lock_tbl1
+ROLLBACK;
+RESET ROLE;
+-- succeed with only SELECT permissions and ACCESS EXCLUSIVE mode
+GRANT SELECT ON TABLE lock_tbl1 TO regress_rol_lock1;
+LOCK TABLE ONLY lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+RESET ROLE;
+REVOKE SELECT ON TABLE lock_tbl1 FROM regress_rol_lock1;
 -- Child tables are locked without granting explicit permission to do so as
 -- long as we have permission to lock the parent.
 GRANT UPDATE ON TABLE lock_tbl1 TO regress_rol_lock1;
diff --git a/src/test/regress/sql/lock.sql b/src/test/regress/sql/lock.sql
index b88488c6d0..df4be3147b 100644
--- a/src/test/regress/sql/lock.sql
+++ b/src/test/regress/sql/lock.sql
@@ -47,6 +47,24 @@ LOCK TABLE lock_tbl1 IN EXCLUSIVE MODE NOWAIT;
 LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE NOWAIT;
 ROLLBACK;
 
+-- Try using WAIT ONLY along with valid options.
+BEGIN TRANSACTION;
+LOCK TABLE lock_tbl1 IN ACCESS SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ROW SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ROW EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE UPDATE EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN SHARE ROW EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN EXCLUSIVE MODE WAIT ONLY;
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+ROLLBACK;
+
+-- WAIT ONLY is allowed outside a transaction
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+
+-- NOWAIT + WAIT ONLY is not supported (yet?)
+LOCK TABLE lock_tbl1 IN ACCESS EXCLUSIVE MODE NOWAIT WAIT ONLY;
+
 -- Verify that we can lock views.
 BEGIN TRANSACTION;
 LOCK TABLE lock_view1 IN EXCLUSIVE MODE;
@@ -104,9 +122,23 @@ ROLLBACK;
 CREATE TABLE lock_tbl2 (b BIGINT) INHERITS (lock_tbl1);
 CREATE TABLE lock_tbl3 () INHERITS (lock_tbl2);
 BEGIN TRANSACTION;
+LOCK TABLE lock_tbl1 * IN ACCESS EXCLUSIVE MODE WAIT ONLY;
 LOCK TABLE lock_tbl1 * IN ACCESS EXCLUSIVE MODE;
 ROLLBACK;
 
+-- WAIT ONLY requires SELECT permissions regardless of lock mode
+-- fail without permissions
+SET ROLE regress_rol_lock1;
+BEGIN;
+LOCK TABLE ONLY lock_tbl1 IN ACCESS SHARE MODE WAIT ONLY;
+ROLLBACK;
+RESET ROLE;
+-- succeed with only SELECT permissions and ACCESS EXCLUSIVE mode
+GRANT SELECT ON TABLE lock_tbl1 TO regress_rol_lock1;
+LOCK TABLE ONLY lock_tbl1 IN ACCESS EXCLUSIVE MODE WAIT ONLY;
+RESET ROLE;
+REVOKE SELECT ON TABLE lock_tbl1 FROM regress_rol_lock1;
+
 -- Child tables are locked without granting explicit permission to do so as
 -- long as we have permission to lock the parent.
 GRANT UPDATE ON TABLE lock_tbl1 TO regress_rol_lock1;
-- 
2.25.1

