Merge #44015

44015: sql/opt/optbuilder: propagate FOR UPDATE locking in scope to ScanExprs r=nvanbenschoten a=nvanbenschoten

Relates to #40205.

This change updates optbuilder to propagate row-level locking modes through the `*scope` object 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 #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.

Co-authored-by: Nathan VanBenschoten <[email protected]>

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -25,25 +25,36 @@ 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.
-
-query I
+query error pgcode 42P01 relation "a" in FOR UPDATE clause not found in FROM clause
 SELECT 1 FOR UPDATE OF a
-----
-1
 
-query I
+query error pgcode 42P01 relation "a" in FOR SHARE clause not found in FROM clause
 SELECT 1 FOR SHARE OF a, b
-----
-1
 
-query I
+query error pgcode 42P01 relation "a" in FOR UPDATE clause not found in FROM clause
 SELECT 1 FOR UPDATE OF a FOR SHARE OF b, c FOR NO KEY UPDATE OF d FOR KEY SHARE OF e, f
+
+query I
+SELECT 1 FROM
+    (SELECT 1) a,
+    (SELECT 1) b,
+    (SELECT 1) c,
+    (SELECT 1) d,
+    (SELECT 1) e,
+    (SELECT 1) f
+FOR UPDATE OF a FOR SHARE OF b, c FOR NO KEY UPDATE OF d FOR KEY SHARE OF e, f
 ----
 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 +121,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

@@ -248,8 +248,13 @@ func (b *Builder) buildStmt(
 		}
 	}
 
-	// NB: The case statements are sorted lexicographically.
 	switch stmt := stmt.(type) {
+	case *tree.Select:
+		return b.buildSelect(stmt, noRowLocking, desiredTypes, inScope)
+
+	case *tree.ParenSelect:
+		return b.buildSelect(stmt.Select, noRowLocking, desiredTypes, inScope)
+
 	case *tree.Delete:
 		return b.processWiths(stmt.With, inScope, func(inScope *scope) *scope {
 			return b.buildDelete(stmt, inScope)
@@ -260,33 +265,6 @@ func (b *Builder) buildStmt(
 			return b.buildInsert(stmt, inScope)
 		})
 
-	case *tree.ParenSelect:
-		return b.processWiths(stmt.Select.With, inScope, func(inScope *scope) *scope {
-			return b.buildSelect(stmt.Select, desiredTypes, inScope)
-		})
-
-	case *tree.Select:
-		rStmt := stmt
-		with := stmt.With
-		wrapped := stmt.Select
-		for s, ok := wrapped.(*tree.ParenSelect); ok; s, ok = wrapped.(*tree.ParenSelect) {
-			stmt = s.Select
-			if stmt.With != nil {
-				if with != nil {
-					// (WITH ... (WITH ...))
-					// Currently we are unable to nest the scopes inside ParenSelect so we
-					// must refuse the syntax so that the query does not get invalid results.
-					panic(unimplemented.NewWithIssue(24303, "multiple WITH clauses in parentheses"))
-				}
-				with = s.Select.With
-			}
-			wrapped = stmt.Select
-		}
-
-		return b.processWiths(with, inScope, func(inScope *scope) *scope {
-			return b.buildSelect(rStmt, desiredTypes, inScope)
-		})
-
 	case *tree.Update:
 		return b.processWiths(stmt.With, inScope, func(inScope *scope) *scope {
 			return b.buildUpdate(stmt, inScope)

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,
 	)