Merge #129768

129768: opt: propagate extra columns for lock operators r=DrewKimball a=DrewKimball

A lock operator needs access to all primary key columns of the relations it locks. Therefore, all primary key columns should be propagated by a `SELECT` clause with locking even if they weren't explicitly projected. Previously, only explicitly projected columns were considered when handling a nested scope, as for a subquery. This could lead to an internal error when the lock operator did not have access to the expected columns. This commit fixes the bug by propgating all primary key columns, not just explicit ones.

Fixes #129647

Release note (bug fix): Fixed a bug that could cause an internal error if a table with an implicit (rowid) primary key was locked from within a subquery, like this:
```
SELECT * FROM (SELECT * FROM foo WHERE x = 2) FOR UPDATE;
```

Co-authored-by: Drew Kimball <[email protected]>

pkg/sql/logictest/testdata/logic_test/dangerous_statements

@@ -40,13 +40,9 @@ statement error rejected.*: SELECT FOR UPDATE without WHERE or LIMIT clause
 statement ok
 (SELECT * FROM foo WHERE x = 2) FOR UPDATE
 
-# Skipped due to https://github.com/cockroachdb/cockroach/issues/129647.
-skipif config weak-iso-level-configs
 statement ok
 SELECT * FROM (SELECT * FROM foo WHERE x = 2) FOR UPDATE
 
-# Skipped due to https://github.com/cockroachdb/cockroach/issues/129647.
-skipif config weak-iso-level-configs
 statement ok
 SELECT * FROM (SELECT * FROM (SELECT * FROM foo) WHERE x = 2) FOR UPDATE
 
@@ -59,8 +55,6 @@ SELECT * FROM (SELECT * FROM foo WHERE x = 2 FOR UPDATE) m, (SELECT * FROM foo)
 statement error rejected.*: SELECT FOR SHARE without WHERE or LIMIT clause
 SELECT * FROM (SELECT * FROM foo FOR SHARE) m, (SELECT * FROM foo) n WHERE m.x = n.x
 
-# Skipped due to https://github.com/cockroachdb/cockroach/issues/129647.
-skipif config weak-iso-level-configs
 statement ok
 SELECT * FROM (SELECT * FROM (SELECT * FROM foo) WHERE x > 1) WHERE x > 2 FOR UPDATE
 

pkg/sql/opt/optbuilder/locking.go

@@ -275,11 +275,13 @@ func (b *Builder) analyzeLockArgs(
 	lockScope = inScope.push()
 	lockScope.cols = make([]scopeColumn, 0, pkCols.Len())
 
-	for i := range inScope.cols {
-		if pkCols.Contains(inScope.cols[i].id) {
-			lockScope.appendColumn(&inScope.cols[i])
+	// Make sure to check extra columns, since the primary key columns may not
+	// have been explicitly projected.
+	inScope.forEachColWithExtras(func(col *scopeColumn) {
+		if pkCols.Contains(col.id) {
+			lockScope.appendColumn(col)
 		}
-	}
+	})
 	return lockScope
 }
 

pkg/sql/opt/optbuilder/scope.go

@@ -568,6 +568,17 @@ func (s *scope) colList() opt.ColList {
 	return colList
 }
 
+// forEachColWithExtras applies the given function to every column in the scope,
+// including extra columns.
+func (s *scope) forEachColWithExtras(fn func(col *scopeColumn)) {
+	for i := range s.cols {
+		fn(&s.cols[i])
+	}
+	for i := range s.extraCols {
+		fn(&s.extraCols[i])
+	}
+}
+
 // hasSameColumns returns true if this scope has the same columns
 // as the other scope.
 //