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

This commit introduces a number of validation checks into the optbuilder
to ensure that FOR UPDATE locking clauses are only being used in places
where they're allowed. Most of this was based off of Postgres, which
disallows the same set of operations to be used with FOR [KEY] UPDATE/SHARE
locking.

Ref: https://github.com/postgres/postgres/blob/1a4a0329650b0545a54afb3c317aa289fd817f8a/src/backend/parser/analyze.c#L2691

Release note (sql change): Invalid usages of FOR UPDATE locking clauses
are now rejected by the SQL optimizer.

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -87,8 +87,53 @@ 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)
@@ -780,40 +780,17 @@ func (b *Builder) buildSelect(
 		}
 	}
 
-	for _, lockItem := range locking {
-		switch lockItem.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 lockItem.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"))
-		}
-	}
-
 	// 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:
@@ -852,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)
@@ -909,6 +890,8 @@ func (b *Builder) buildSelectClause(
 			outScope = b.buildDistinctOn(projectionsScope.distinctOnCols, outScope)
 		}
 	}
+
+	b.validateLockingForSelectClause(sel, locking, fromScope)
 	return outScope
 }
 
@@ -1102,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/sem/tree/select.go

@@ -985,6 +985,7 @@ type LockingStrength byte
 
 const (
 	// ForNone represents the default - no for statement at all.
+	// LockingItem AST nodes are never created with this strength.
 	ForNone LockingStrength = iota
 	// ForUpdate represents FOR UPDATE.
 	ForUpdate
@@ -996,18 +997,23 @@ const (
 	ForKeyShare
 )
 
+var lockingStrengthName = [...]string{
+	ForNone:        "",
+	ForUpdate:      "FOR UPDATE",
+	ForNoKeyUpdate: "FOR NO KEY UPDATE",
+	ForShare:       "FOR SHARE",
+	ForKeyShare:    "FOR KEY SHARE",
+}
+
+func (s LockingStrength) String() string {
+	return lockingStrengthName[s]
+}
+
 // Format implements the NodeFormatter interface.
-func (f LockingStrength) Format(ctx *FmtCtx) {
-	switch f {
-	case ForNone:
-	case ForUpdate:
-		ctx.WriteString(" FOR UPDATE")
-	case ForNoKeyUpdate:
-		ctx.WriteString(" FOR NO KEY UPDATE")
-	case ForShare:
-		ctx.WriteString(" FOR SHARE")
-	case ForKeyShare:
-		ctx.WriteString(" FOR KEY SHARE")
+func (s LockingStrength) Format(ctx *FmtCtx) {
+	if s != ForNone {
+		ctx.WriteString(" ")
+		ctx.WriteString(s.String())
 	}
 }
 
@@ -1027,13 +1033,20 @@ const (
 	LockWaitError
 )
 
+var lockingWaitPolicyName = [...]string{
+	LockWaitBlock: "",
+	LockWaitSkip:  "SKIP LOCKED",
+	LockWaitError: "NOWAIT",
+}
+
+func (p LockingWaitPolicy) String() string {
+	return lockingWaitPolicyName[p]
+}
+
 // Format implements the NodeFormatter interface.
-func (f LockingWaitPolicy) Format(ctx *FmtCtx) {
-	switch f {
-	case LockWaitBlock:
-	case LockWaitSkip:
-		ctx.WriteString(" SKIP LOCKED")
-	case LockWaitError:
-		ctx.WriteString(" NOWAIT")
+func (p LockingWaitPolicy) Format(ctx *FmtCtx) {
+	if p != LockWaitBlock {
+		ctx.WriteString(" ")
+		ctx.WriteString(p.String())
 	}
 }