sql/opt/optbuilder: propagate FOR UPDATE locking in scope to ScanExprs

Relates to cockroachdb#40205.

This change updates optbuilder to propagate row-level locking modes on the stack
to achieve the same locking clause scoping that Postgres contains. The change
propagates the locking mode all the way down to the ScanExpr level, where a
single locking mode becomes an optional property of a Scan.

With proper scoping rules, the change also improves upon the locking validation
introduced in cockroachdb#43887.

This completes the SELECT FOR [KEY] SHARE/UPDATE optbuilder interactions. The
next step will be to hook the ScanExpr locking property into the execbuilder.

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

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -44,6 +44,15 @@ SELECT 1 FOR UPDATE OF a FOR SHARE OF b, c FOR NO KEY UPDATE OF d FOR KEY SHARE
 ----
 1
 
+# However, we do mirror Postgres in that we require FOR UPDATE targets to be
+# unqualified names and reject anything else.
+
+query error pgcode 42601 FOR UPDATE must specify unqualified relation names
+SELECT 1 FOR UPDATE OF public.a
+
+query error pgcode 42601 FOR UPDATE must specify unqualified relation names
+SELECT 1 FOR UPDATE OF db.public.a
+
 # 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.
@@ -110,30 +119,99 @@ SELECT 1 FOR READ ONLY
 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 UNION/INTERSECT/EXCEPT
+SELECT * FROM (SELECT 1 UNION SELECT 1) a 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 VALUES
+SELECT * FROM (VALUES (1)) a 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 DISTINCT clause
+SELECT * FROM (SELECT DISTINCT 1) a 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 GROUP BY clause
+SELECT * FROM (SELECT 1 GROUP BY 1) a 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 HAVING clause
+SELECT * FROM (SELECT 1 HAVING TRUE) a 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 aggregate functions
+SELECT * FROM (SELECT count(1)) a 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 window functions
+SELECT * FROM (SELECT count(1) OVER ()) a 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
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with set-returning functions in the target list
+SELECT * FROM (SELECT generate_series(1, 2)) a FOR UPDATE
+
 # Set-returning functions in the from list are allowed.
 query I
 SELECT * FROM generate_series(1, 2) FOR UPDATE
 ----
 1
 2
+
+query I
+SELECT * FROM (SELECT * FROM generate_series(1, 2)) a FOR UPDATE
+----
+1
+2
+
+# Use of SELECT FOR UPDATE/SHARE requires UPDATE privileges, not just SELECT privileges.
+
+statement ok
+CREATE TABLE t (k INT PRIMARY KEY, v int)
+
+statement ok
+GRANT GRANT on t to testuser
+
+user testuser
+
+statement error pgcode 42501 user testuser does not have SELECT privilege on relation t
+SELECT * FROM t
+
+statement ok
+GRANT SELECT ON t TO testuser
+
+statement ok
+SELECT * FROM t
+
+statement error pgcode 42501 user testuser does not have UPDATE privilege on relation t
+SELECT * FROM t FOR UPDATE
+
+statement error pgcode 42501 user testuser does not have UPDATE privilege on relation t
+SELECT * FROM t FOR SHARE
+
+statement ok
+GRANT UPDATE ON t TO testuser
+
+statement ok
+SELECT * FROM t FOR UPDATE
+
+statement ok
+SELECT * FROM t FOR SHARE
+
+user root
+
+statement ok
+DROP TABLE t

pkg/sql/opt/memo/expr_format.go

@@ -360,6 +360,29 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 				tp.Childf("flags: force-index=%s%s", idx.Name(), dir)
 			}
 		}
+		if t.Locking != nil {
+			strength := ""
+			switch t.Locking.Strength {
+			case tree.ForNone:
+			case tree.ForKeyShare:
+				strength = "for-key-share"
+			case tree.ForShare:
+				strength = "for-share"
+			case tree.ForNoKeyUpdate:
+				strength = "for-no-key-update"
+			case tree.ForUpdate:
+				strength = "for-update"
+			}
+			wait := ""
+			switch t.Locking.WaitPolicy {
+			case tree.LockWaitBlock:
+			case tree.LockWaitSkip:
+				wait = ",skip-locked"
+			case tree.LockWaitError:
+				wait = ",nowait"
+			}
+			tp.Childf("locking: %s%s", strength, wait)
+		}
 
 	case *LookupJoinExpr:
 		if !t.Flags.Empty() {

pkg/sql/opt/memo/interner.go

@@ -311,17 +311,24 @@ func (h *hasher) HashFloat64(val float64) {
 	h.hash *= prime64
 }
 
+func (h *hasher) HashRune(val rune) {
+	h.hash ^= internHash(val)
+	h.hash *= prime64
+}
+
 func (h *hasher) HashString(val string) {
 	for _, c := range val {
-		h.hash ^= internHash(c)
-		h.hash *= prime64
+		h.HashRune(c)
 	}
 }
 
+func (h *hasher) HashByte(val byte) {
+	h.HashRune(rune(val))
+}
+
 func (h *hasher) HashBytes(val []byte) {
 	for _, c := range val {
-		h.hash ^= internHash(c)
-		h.hash *= prime64
+		h.HashByte(c)
 	}
 }
 
@@ -540,6 +547,13 @@ func (h *hasher) HashPhysProps(val *physical.Required) {
 	h.HashFloat64(val.LimitHint)
 }
 
+func (h *hasher) HashLockingItem(val *tree.LockingItem) {
+	if val != nil {
+		h.HashByte(byte(val.Strength))
+		h.HashByte(byte(val.WaitPolicy))
+	}
+}
+
 func (h *hasher) HashRelExpr(val RelExpr) {
 	h.HashUint64(uint64(reflect.ValueOf(val).Pointer()))
 }
@@ -646,10 +660,18 @@ func (h *hasher) IsFloat64Equal(l, r float64) bool {
 	return math.Float64bits(l) == math.Float64bits(r)
 }
 
+func (h *hasher) IsRuneEqual(l, r rune) bool {
+	return l == r
+}
+
 func (h *hasher) IsStringEqual(l, r string) bool {
 	return l == r
 }
 
+func (h *hasher) IsByteEqual(l, r byte) bool {
+	return l == r
+}
+
 func (h *hasher) IsBytesEqual(l, r []byte) bool {
 	return bytes.Equal(l, r)
 }
@@ -854,6 +876,13 @@ func (h *hasher) IsPhysPropsEqual(l, r *physical.Required) bool {
 	return l.Equals(r)
 }
 
+func (h *hasher) IsLockingItemEqual(l, r *tree.LockingItem) bool {
+	if l == nil || r == nil {
+		return l == r
+	}
+	return l.Strength == r.Strength && l.WaitPolicy == r.WaitPolicy
+}
+
 func (h *hasher) IsPointerEqual(l, r unsafe.Pointer) bool {
 	return l == r
 }

pkg/sql/opt/memo/interner_test.go

@@ -177,13 +177,27 @@ func TestInterner(t *testing.T) {
 			{val1: float64(0), val2: math.Copysign(0, -1), equal: false},
 		}},
 
+		{hashFn: in.hasher.HashRune, eqFn: in.hasher.IsRuneEqual, variations: []testVariation{
+			{val1: rune(0), val2: rune(0), equal: true},
+			{val1: rune('a'), val2: rune('b'), equal: false},
+			{val1: rune('a'), val2: rune('A'), equal: false},
+			{val1: rune('🐛'), val2: rune('🐛'), equal: true},
+		}},
+
 		{hashFn: in.hasher.HashString, eqFn: in.hasher.IsStringEqual, variations: []testVariation{
 			{val1: "", val2: "", equal: true},
 			{val1: "abc", val2: "abcd", equal: false},
 			{val1: "", val2: " ", equal: false},
 			{val1: "the quick brown fox", val2: "the quick brown fox", equal: true},
 		}},
 
+		{hashFn: in.hasher.HashByte, eqFn: in.hasher.IsByteEqual, variations: []testVariation{
+			{val1: byte(0), val2: byte(0), equal: true},
+			{val1: byte('a'), val2: byte('b'), equal: false},
+			{val1: byte('a'), val2: byte('A'), equal: false},
+			{val1: byte('z'), val2: byte('z'), equal: true},
+		}},
+
 		{hashFn: in.hasher.HashBytes, eqFn: in.hasher.IsBytesEqual, variations: []testVariation{
 			{val1: []byte{}, val2: []byte{}, equal: true},
 			{val1: []byte{}, val2: []byte{0}, equal: false},
@@ -412,6 +426,30 @@ func TestInterner(t *testing.T) {
 
 		// PhysProps hash/isEqual methods are tested in TestInternerPhysProps.
 
+		{hashFn: in.hasher.HashLockingItem, eqFn: in.hasher.IsLockingItemEqual, variations: []testVariation{
+			{val1: (*tree.LockingItem)(nil), val2: (*tree.LockingItem)(nil), equal: true},
+			{
+				val1:  (*tree.LockingItem)(nil),
+				val2:  &tree.LockingItem{Strength: tree.ForUpdate},
+				equal: false,
+			},
+			{
+				val1:  &tree.LockingItem{Strength: tree.ForShare},
+				val2:  &tree.LockingItem{Strength: tree.ForUpdate},
+				equal: false,
+			},
+			{
+				val1:  &tree.LockingItem{WaitPolicy: tree.LockWaitSkip},
+				val2:  &tree.LockingItem{WaitPolicy: tree.LockWaitError},
+				equal: false,
+			},
+			{
+				val1:  &tree.LockingItem{Strength: tree.ForUpdate, WaitPolicy: tree.LockWaitError},
+				val2:  &tree.LockingItem{Strength: tree.ForUpdate, WaitPolicy: tree.LockWaitError},
+				equal: true,
+			},
+		}},
+
 		{hashFn: in.hasher.HashRelExpr, eqFn: in.hasher.IsRelExprEqual, variations: []testVariation{
 			{val1: (*ScanExpr)(nil), val2: (*ScanExpr)(nil), equal: true},
 			{val1: scanNode, val2: scanNode, equal: true},

pkg/sql/opt/ops/relational.opt

@@ -65,6 +65,14 @@ define ScanPrivate {
     # Flags modify how the table is scanned, such as which index is used to scan.
     Flags ScanFlags
 
+    # Locking represents the row-level locking mode of the Scan. Most scans
+    # leave this unset (Strength = ForNone), which indicates that no row-level
+    # locking will be performed while scanning the table. Stronger locking modes
+    # are used by SELECT .. FOR [KEY] UPDATE/SHARE statements and by the initial
+    # row retrieval of DELETE and UPDATE statements. The locking item's Targets
+    # list will always be empty when part of a ScanPrivate.
+    Locking LockingItem
+
     # PartitionConstrainedScan records whether or not we were able to use partitions
     # to constrain the lookup spans further. This flag is used to record telemetry
     # about how often this optimization is getting applied.

pkg/sql/opt/optbuilder/builder.go

@@ -250,10 +250,10 @@ func (b *Builder) buildStmt(
 
 	switch stmt := stmt.(type) {
 	case *tree.Select:
-		return b.buildSelect(stmt, desiredTypes, inScope)
+		return b.buildSelect(stmt, noRowLocking, desiredTypes, inScope)
 
 	case *tree.ParenSelect:
-		return b.buildSelect(stmt.Select, desiredTypes, inScope)
+		return b.buildSelect(stmt.Select, noRowLocking, desiredTypes, inScope)
 
 	case *tree.Delete:
 		return b.processWiths(stmt.With, inScope, func(inScope *scope) *scope {

pkg/sql/opt/optbuilder/insert.go

@@ -664,6 +664,7 @@ func (mb *mutationBuilder) buildInputForDoNothing(inScope *scope, onConflict *tr
 			mb.b.addTable(mb.tab, &mb.alias),
 			nil, /* ordinals */
 			nil, /* indexFlags */
+			noRowLocking,
 			excludeMutations,
 			inScope,
 		)
@@ -746,6 +747,7 @@ func (mb *mutationBuilder) buildInputForUpsert(
 		mb.b.addTable(mb.tab, &mb.alias),
 		nil, /* ordinals */
 		nil, /* indexFlags */
+		noRowLocking,
 		includeMutations,
 		inScope,
 	)

pkg/sql/opt/optbuilder/join.go

@@ -29,8 +29,10 @@ import (
 //
 // See Builder.buildStmt for a description of the remaining input and
 // return values.
-func (b *Builder) buildJoin(join *tree.JoinTableExpr, inScope *scope) (outScope *scope) {
-	leftScope := b.buildDataSource(join.Left, nil /* indexFlags */, inScope)
+func (b *Builder) buildJoin(
+	join *tree.JoinTableExpr, locking lockingSpec, inScope *scope,
+) (outScope *scope) {
+	leftScope := b.buildDataSource(join.Left, nil /* indexFlags */, locking, inScope)
 
 	isLateral := false
 	inScopeRight := inScope
@@ -43,7 +45,7 @@ func (b *Builder) buildJoin(join *tree.JoinTableExpr, inScope *scope) (outScope
 		inScopeRight = leftScope
 	}
 
-	rightScope := b.buildDataSource(join.Right, nil /* indexFlags */, inScopeRight)
+	rightScope := b.buildDataSource(join.Right, nil /* indexFlags */, locking, inScopeRight)
 
 	// Check that the same table name is not used on both sides.
 	b.validateJoinTableNames(leftScope, rightScope)