pkg/sql/opt: validate correct use of FOR UPDATE locking clauses

docs/generated/sql/bnf/select_stmt.bnf

@@ -1,3 +1,3 @@
 select_stmt ::=
-	( simple_select locking_clause | select_clause sort_clause locking_clause | select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause sort_clause locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) locking_clause )
+	( select_clause sort_clause | select_clause ( sort_clause |  ) for_locking_clause opt_select_limit | select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) opt_for_locking_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause sort_clause | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) for_locking_clause opt_select_limit | ( 'WITH' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) | 'WITH' 'RECURSIVE' ( ( common_table_expr ) ( ( ',' common_table_expr ) )* ) ) select_clause ( sort_clause |  ) ( limit_clause offset_clause | offset_clause limit_clause | limit_clause | offset_clause ) opt_for_locking_clause )
 

docs/generated/sql/bnf/stmt_block.bnf

@@ -439,12 +439,14 @@ scrub_database_stmt ::=
 	'EXPERIMENTAL' 'SCRUB' 'DATABASE' database_name opt_as_of_clause
 
 select_no_parens ::=
-	simple_select locking_clause
-	| select_clause sort_clause locking_clause
-	| select_clause opt_sort_clause select_limit locking_clause
-	| with_clause select_clause locking_clause
-	| with_clause select_clause sort_clause locking_clause
-	| with_clause select_clause opt_sort_clause select_limit locking_clause
+	simple_select
+	| select_clause sort_clause
+	| select_clause opt_sort_clause for_locking_clause opt_select_limit
+	| select_clause opt_sort_clause select_limit opt_for_locking_clause
+	| with_clause select_clause
+	| with_clause select_clause sort_clause
+	| with_clause select_clause opt_sort_clause for_locking_clause opt_select_limit
+	| with_clause select_clause opt_sort_clause select_limit opt_for_locking_clause
 
 select_with_parens ::=
 	'(' select_no_parens ')'
@@ -1172,19 +1174,28 @@ simple_select ::=
 	| table_clause
 	| set_operation
 
-locking_clause ::=
-	for_locking_strength opt_locked_rels opt_nowait_or_skip
-
 select_clause ::=
 	simple_select
 	| select_with_parens
 
+for_locking_clause ::=
+	for_locking_items
+	| 'FOR' 'READ' 'ONLY'
+
+opt_select_limit ::=
+	select_limit
+	| 
+
 select_limit ::=
 	limit_clause offset_clause
 	| offset_clause limit_clause
 	| limit_clause
 	| offset_clause
 
+opt_for_locking_clause ::=
+	for_locking_clause
+	| 
+
 set_rest_more ::=
 	generic_set
 
@@ -1574,18 +1585,8 @@ set_operation ::=
 	| select_clause 'INTERSECT' all_or_distinct select_clause
 	| select_clause 'EXCEPT' all_or_distinct select_clause
 
-for_locking_strength ::=
-	'FOR' 'UPDATE'
-	| 'FOR' 'NO' 'KEY' 'UPDATE'
-	| 'FOR' 'SHARE'
-	| 'FOR' 'KEY' 'SHARE'
-
-opt_locked_rels ::=
-	'OF' table_name_list
-
-opt_nowait_or_skip ::=
-	'SKIP' 'LOCKED'
-	| 'NOWAIT'
+for_locking_items ::=
+	( for_locking_item ) ( ( for_locking_item ) )*
 
 offset_clause ::=
 	'OFFSET' a_expr
@@ -1939,6 +1940,9 @@ all_or_distinct ::=
 	| 'DISTINCT'
 	| 
 
+for_locking_item ::=
+	for_locking_strength opt_locked_rels opt_nowait_or_skip
+
 var_list ::=
 	( var_value ) ( ( ',' var_value ) )*
 
@@ -2204,6 +2208,19 @@ interval_qualifier ::=
 window_definition_list ::=
 	( window_definition ) ( ( ',' window_definition ) )*
 
+for_locking_strength ::=
+	'FOR' 'UPDATE'
+	| 'FOR' 'NO' 'KEY' 'UPDATE'
+	| 'FOR' 'SHARE'
+	| 'FOR' 'KEY' 'SHARE'
+
+opt_locked_rels ::=
+	'OF' table_name_list
+
+opt_nowait_or_skip ::=
+	'SKIP' 'LOCKED'
+	| 'NOWAIT'
+
 opt_join_hint ::=
 	'HASH'
 	| 'MERGE'

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -20,6 +20,11 @@ SELECT 1 FOR KEY SHARE
 ----
 1
 
+query I
+SELECT 1 FOR UPDATE FOR SHARE FOR NO KEY UPDATE FOR KEY SHARE
+----
+1
+
 # Postgres gives an error if you specify a table that isn't available in the
 # FROM list for the OF ... clause, but since we're ignoring this list anyway,
 # we didn't bother to add that behavior.
@@ -34,6 +39,11 @@ SELECT 1 FOR SHARE OF a, b
 ----
 1
 
+query I
+SELECT 1 FOR UPDATE OF a FOR SHARE OF b, c FOR NO KEY UPDATE OF d FOR KEY SHARE OF e, f
+----
+1
+
 # We can't support SKIP LOCKED or NOWAIT, since they would actually behave
 # differently - NOWAIT returns an error to the client instead of blocking,
 # and SKIP LOCKED returns an inconsistent view.
@@ -53,6 +63,12 @@ SELECT 1 FOR KEY SHARE SKIP LOCKED
 query error unimplemented: SKIP LOCKED lock wait policy is not supported
 SELECT 1 FOR UPDATE OF a SKIP LOCKED
 
+query error unimplemented: SKIP LOCKED lock wait policy is not supported
+SELECT 1 FOR UPDATE OF a SKIP LOCKED FOR NO KEY UPDATE OF b SKIP LOCKED
+
+query error unimplemented: SKIP LOCKED lock wait policy is not supported
+SELECT 1 FOR UPDATE OF a SKIP LOCKED FOR NO KEY UPDATE OF b NOWAIT
+
 query error unimplemented: NOWAIT lock wait policy is not supported
 SELECT 1 FOR UPDATE NOWAIT
 
@@ -67,3 +83,57 @@ SELECT 1 FOR KEY SHARE NOWAIT
 
 query error unimplemented: NOWAIT lock wait policy is not supported
 SELECT 1 FOR UPDATE OF a NOWAIT
+
+query error unimplemented: NOWAIT lock wait policy is not supported
+SELECT 1 FOR UPDATE OF a NOWAIT FOR NO KEY UPDATE OF b NOWAIT
+
+# Locking clauses both inside and outside of parenthesis are handled correctly.
+
+query error unimplemented: SKIP LOCKED lock wait policy is not supported
+((SELECT 1)) FOR UPDATE SKIP LOCKED
+
+query error unimplemented: SKIP LOCKED lock wait policy is not supported
+((SELECT 1) FOR UPDATE SKIP LOCKED)
+
+query error unimplemented: SKIP LOCKED lock wait policy is not supported
+((SELECT 1 FOR UPDATE SKIP LOCKED))
+
+# FOR READ ONLY is ignored, like in Postgres.
+query I
+SELECT 1 FOR READ ONLY
+----
+1
+
+# Various operations are not supported when locking clauses are provided.
+# FeatureNotSupported errors are thrown for each of them.
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with UNION/INTERSECT/EXCEPT
+SELECT 1 UNION SELECT 1 FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with VALUES
+VALUES (1) FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with DISTINCT clause
+SELECT DISTINCT 1 FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with GROUP BY clause
+SELECT 1 GROUP BY 1 FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with HAVING clause
+SELECT 1 HAVING TRUE FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with aggregate functions
+SELECT count(1) FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with window functions
+SELECT count(1) OVER () FOR UPDATE
+
+statement error pgcode 0A000 FOR UPDATE is not allowed with set-returning functions in the target list
+SELECT generate_series(1, 2) FOR UPDATE
+
+# Set-returning functions in the from list are allowed.
+query I
+SELECT * FROM generate_series(1, 2) FOR UPDATE
+----
+1
+2

pkg/sql/opt/optbuilder/select.go

@@ -719,7 +719,7 @@ func (b *Builder) buildSelectStmt(
 		return b.buildStmt(stmt.Select, desiredTypes, inScope)
 
 	case *tree.SelectClause:
-		return b.buildSelectClause(stmt, nil /* orderBy */, desiredTypes, inScope)
+		return b.buildSelectClause(stmt, nil /* orderBy */, nil /* locking */, desiredTypes, inScope)
 
 	case *tree.UnionClause:
 		return b.buildUnionClause(stmt, desiredTypes, inScope)
@@ -754,30 +754,7 @@ func (b *Builder) buildSelect(
 	wrapped := stmt.Select
 	orderBy := stmt.OrderBy
 	limit := stmt.Limit
-	forLocked := stmt.ForLocked
-
-	switch forLocked.Strength {
-	case tree.ForNone:
-	case tree.ForUpdate:
-	case tree.ForNoKeyUpdate:
-	case tree.ForShare:
-	case tree.ForKeyShare:
-		// CockroachDB treats all of the FOR LOCKED modes as no-ops. Since all
-		// transactions are serializable in CockroachDB, clients can't observe
-		// whether or not FOR UPDATE (or any of the other weaker modes) actually
-		// created a lock. This behavior may improve as the transaction model gains
-		// more capabilities.
-	}
-
-	switch forLocked.WaitPolicy {
-	case tree.LockWaitBlock:
-	case tree.LockWaitSkip:
-		panic(unimplementedWithIssueDetailf(40476, "",
-			"SKIP LOCKED lock wait policy is not supported"))
-	case tree.LockWaitError:
-		panic(unimplementedWithIssueDetailf(40476, "",
-			"NOWAIT lock wait policy is not supported"))
-	}
+	locking := stmt.Locking
 
 	for s, ok := wrapped.(*tree.ParenSelect); ok; s, ok = wrapped.(*tree.ParenSelect) {
 		stmt = s.Select
@@ -798,17 +775,22 @@ func (b *Builder) buildSelect(
 			}
 			limit = stmt.Limit
 		}
+		if stmt.Locking != nil {
+			locking = append(locking, stmt.Locking...)
+		}
 	}
 
 	// NB: The case statements are sorted lexicographically.
 	switch t := stmt.Select.(type) {
 	case *tree.SelectClause:
-		outScope = b.buildSelectClause(t, orderBy, desiredTypes, inScope)
+		outScope = b.buildSelectClause(t, orderBy, locking, desiredTypes, inScope)
 
 	case *tree.UnionClause:
+		b.rejectIfLocking(locking, "UNION/INTERSECT/EXCEPT")
 		outScope = b.buildUnionClause(t, desiredTypes, inScope)
 
 	case *tree.ValuesClause:
+		b.rejectIfLocking(locking, "VALUES")
 		outScope = b.buildValuesClause(t, desiredTypes, inScope)
 
 	default:
@@ -847,7 +829,11 @@ func (b *Builder) buildSelect(
 // See Builder.buildStmt for a description of the remaining input and
 // return values.
 func (b *Builder) buildSelectClause(
-	sel *tree.SelectClause, orderBy tree.OrderBy, desiredTypes []*types.T, inScope *scope,
+	sel *tree.SelectClause,
+	orderBy tree.OrderBy,
+	locking tree.LockingClause,
+	desiredTypes []*types.T,
+	inScope *scope,
 ) (outScope *scope) {
 	fromScope := b.buildFrom(sel.From, inScope)
 	b.processWindowDefs(sel, fromScope)
@@ -904,6 +890,8 @@ func (b *Builder) buildSelectClause(
 			outScope = b.buildDistinctOn(projectionsScope.distinctOnCols, outScope)
 		}
 	}
+
+	b.validateLockingForSelectClause(sel, locking, fromScope)
 	return outScope
 }
 
@@ -1097,3 +1085,80 @@ func (b *Builder) validateAsOf(asOf tree.AsOfClause) {
 			"cannot specify AS OF SYSTEM TIME with different timestamps"))
 	}
 }
+
+// validateLockingForSelectClause checks for operations that are not supported
+// with FOR [KEY] UPDATE/SHARE. If a locking clause was specified with the
+// select and an incompatible operation is in use, a locking error is raised.
+// The method also validates that only supported locking modes are used.
+func (b *Builder) validateLockingForSelectClause(
+	sel *tree.SelectClause, locking tree.LockingClause, scope *scope,
+) {
+	if len(locking) == 0 {
+		return
+	}
+
+	first := locking[0]
+	switch {
+	case sel.Distinct:
+		b.raiseLockingError(first, "DISTINCT clause")
+
+	case sel.GroupBy != nil:
+		b.raiseLockingError(first, "GROUP BY clause")
+
+	case sel.Having != nil:
+		b.raiseLockingError(first, "HAVING clause")
+
+	case scope.groupby != nil && scope.groupby.hasAggregates():
+		b.raiseLockingError(first, "aggregate functions")
+
+	case len(scope.windows) != 0:
+		b.raiseLockingError(first, "window functions")
+
+	case len(scope.srfs) != 0:
+		b.raiseLockingError(first, "set-returning functions in the target list")
+	}
+
+	for _, li := range locking {
+		// Validate locking strength.
+		switch li.Strength {
+		case tree.ForNone:
+			// AST nodes should not be created with this locking strength.
+			panic(errors.AssertionFailedf("locking item without strength"))
+		case tree.ForUpdate, tree.ForNoKeyUpdate, tree.ForShare, tree.ForKeyShare:
+			// CockroachDB treats all of the FOR LOCKED modes as no-ops. Since all
+			// transactions are serializable in CockroachDB, clients can't observe
+			// whether or not FOR UPDATE (or any of the other weaker modes) actually
+			// created a lock. This behavior may improve as the transaction model gains
+			// more capabilities.
+		default:
+			panic(errors.AssertionFailedf("unknown locking strength: %s", li.Strength))
+		}
+
+		// Validating locking wait policy.
+		switch li.WaitPolicy {
+		case tree.LockWaitBlock:
+			// Default.
+		case tree.LockWaitSkip:
+			panic(unimplementedWithIssueDetailf(40476, "",
+				"SKIP LOCKED lock wait policy is not supported"))
+		case tree.LockWaitError:
+			panic(unimplementedWithIssueDetailf(40476, "",
+				"NOWAIT lock wait policy is not supported"))
+		default:
+			panic(errors.AssertionFailedf("unknown locking wait policy: %s", li.WaitPolicy))
+		}
+	}
+}
+
+// rejectIfLocking raises a locking error if a locking clause was specified.
+func (b *Builder) rejectIfLocking(locking tree.LockingClause, context string) {
+	if len(locking) == 0 {
+		return
+	}
+	b.raiseLockingError(locking[0], context)
+}
+
+func (b *Builder) raiseLockingError(first *tree.LockingItem, context string) {
+	panic(pgerror.Newf(pgcode.FeatureNotSupported,
+		"%s is not allowed with %s", first.Strength, context))
+}

pkg/sql/parser/parse_test.go

@@ -1071,6 +1071,11 @@ func TestParse(t *testing.T) {
 		{`SELECT 1 FOR NO KEY UPDATE OF a, b`},
 		{`SELECT 1 FOR UPDATE SKIP LOCKED`},
 		{`SELECT 1 FOR NO KEY UPDATE OF a, b NOWAIT`},
+		{`SELECT 1 ORDER BY 1 FOR UPDATE`},
+		{`SELECT 1 LIMIT 1 FOR UPDATE`},
+		{`SELECT 1 ORDER BY 1 LIMIT 1 FOR UPDATE`},
+		{`SELECT 1 FOR UPDATE FOR UPDATE`},
+		{`SELECT 1 FOR SHARE OF a FOR KEY SHARE SKIP LOCKED`},
 
 		{`TABLE a`}, // Shorthand for: SELECT * FROM a; used e.g. in CREATE VIEW v AS TABLE t
 		{`EXPLAIN TABLE a`},
@@ -1786,6 +1791,11 @@ func TestParse2(t *testing.T) {
 		{`SELECT ROW(1)`, `SELECT (1,)`},
 		{`SELECT (ROW(1) AS a)`, `SELECT ((1,) AS a)`},
 
+		{`SELECT 1 ORDER BY 1 FOR UPDATE LIMIT 1`,
+			`SELECT 1 ORDER BY 1 LIMIT 1 FOR UPDATE`},
+		// FOR READ ONLY is ignored, like in Postgres.
+		{`SELECT 1 FOR READ ONLY`, `SELECT 1`},
+
 		{`SHOW CREATE TABLE t`,
 			`SHOW CREATE t`},
 		{`SHOW CREATE VIEW t`,