optbuilder: refactor semantic analysis of FOR UPDATE

Locking clauses such as FOR UPDATE and FOR SHARE apply to some or all of
the data sources in a query's FROM list, depending on whether they have
targets (FOR UPDATE OF x). Without targets, they always apply within
subqueries in the FROM list. With targets, they apply within subqueries
if the subquery alias matches the target.

Because of this scope-like nature of FOR UPDATE and FOR SHARE, we
implement semantic analysis using a stack of locking items that grow as
we build inner subqueries deeper in the recursive optbuilder calls.

Prior to this change, we only used the stack of locking items during
buildScan, at the very bottom of the recursion. Because of this, calls
to `lockingSpec.filter` could afford to compress the stack into a single
locking item on our way deeper in the recursion.

As part of the upcoming fix for cockroachdb#75457, however, we will need to build a
new Lock operator when popping off locking items after returning from
the recursion. That Lock operator will need some information gathered
from buildScan at the bottom of the recursion.

To support this, we refactor the stack of locking items to be two
stacks: one that tracks all locking items in scope, and a second that
tracks which locking items currently apply. This will allow buildScan to
associate table information with the correct locking item(s), which can
then be used to build Lock operators when popping the locking items.

As a bonus, by using only the applied locking item stack in
`validateLockingInFrom` we can make validation of SELECT FOR UPDATE
queries a little more precise, which allows some queries we were
incorrectly disallowing.

Informs: cockroachdb#57031, cockroachdb#75457

Epic: CRDB-25322

Release note (sql change): Allow FOR UPDATE on some queries that were
previously disallowed. Queries that use the following operations are now
allowed to have FOR UPDATE OF as long as the prohibited operation is in
a subquery not locked by the FOR UPDATE OF:

- UNION
- INTERSECT
- EXCEPT
- DISTINCT
- GROUP BY
- HAVING
- aggregations
- window functions

For example, the following query is now allowed because the subquery
using the prohibited operations is not affected by the FOR UPDATE OF:

```
SELECT * FROM t,
  (SELECT DISTINCT 0, 0 UNION SELECT a, count(*) FROM t GROUP BY a HAVING a > 0) AS u
  FOR UPDATE OF t;
```

This matches PostgreSQL.

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -172,6 +172,14 @@ 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 UNION/INTERSECT/EXCEPT
+SELECT * FROM (SELECT 1 UNION SELECT 1) a, i FOR UPDATE
+
+query II
+SELECT * FROM (SELECT 1 UNION SELECT 1) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with VALUES
 VALUES (1) FOR UPDATE
 
@@ -195,6 +203,16 @@ SELECT * FROM (VALUES (1)) a, i FOR UPDATE
 ----
 1  1
 
+query II
+SELECT * FROM (VALUES (1)) a, i FOR UPDATE OF a
+----
+1  1
+
+query II
+SELECT * FROM (VALUES (1)) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with DISTINCT clause
 SELECT DISTINCT 1 FOR UPDATE
 
@@ -204,6 +222,14 @@ SELECT * FROM (SELECT DISTINCT 1) a FOR UPDATE
 statement error pgcode 0A000 FOR UPDATE is not allowed with DISTINCT clause
 SELECT * FROM (SELECT DISTINCT 1) a, i FOR UPDATE
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with DISTINCT clause
+SELECT * FROM (SELECT DISTINCT 1) a, i FOR UPDATE OF a
+
+query II
+SELECT * FROM (SELECT DISTINCT 1) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with GROUP BY clause
 SELECT 1 GROUP BY 1 FOR UPDATE
 
@@ -213,6 +239,14 @@ SELECT * FROM (SELECT 1 GROUP BY 1) a FOR UPDATE
 statement error pgcode 0A000 FOR UPDATE is not allowed with GROUP BY clause
 SELECT * FROM (SELECT 1 GROUP BY 1) a, i FOR UPDATE
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with GROUP BY clause
+SELECT * FROM (SELECT 1 GROUP BY 1) a, i FOR UPDATE OF a
+
+query II
+SELECT * FROM (SELECT 1 GROUP BY 1) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with HAVING clause
 SELECT 1 HAVING TRUE FOR UPDATE
 
@@ -222,6 +256,14 @@ SELECT * FROM (SELECT 1 HAVING TRUE) a FOR UPDATE
 statement error pgcode 0A000 FOR UPDATE is not allowed with HAVING clause
 SELECT * FROM (SELECT 1 HAVING TRUE) a, i FOR UPDATE
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with HAVING clause
+SELECT * FROM (SELECT 1 HAVING TRUE) a, i FOR UPDATE OF a
+
+query II
+SELECT * FROM (SELECT 1 HAVING TRUE) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with aggregate functions
 SELECT count(1) FOR UPDATE
 
@@ -231,6 +273,14 @@ SELECT * FROM (SELECT count(1)) a FOR UPDATE
 statement error pgcode 0A000 FOR UPDATE is not allowed with aggregate functions
 SELECT * FROM (SELECT count(1)) a, i FOR UPDATE
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with aggregate functions
+SELECT * FROM (SELECT count(1)) a, i FOR UPDATE OF a
+
+query II
+SELECT * FROM (SELECT count(1)) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with window functions
 SELECT count(1) OVER () FOR UPDATE
 
@@ -240,6 +290,14 @@ SELECT * FROM (SELECT count(1) OVER ()) a FOR UPDATE
 statement error pgcode 0A000 FOR UPDATE is not allowed with window functions
 SELECT * FROM (SELECT count(1) OVER ()) a, i FOR UPDATE
 
+statement error pgcode 0A000 FOR UPDATE is not allowed with window functions
+SELECT * FROM (SELECT count(1) OVER ()) a, i FOR UPDATE OF a
+
+query II
+SELECT * FROM (SELECT count(1) OVER ()) a, i FOR UPDATE OF i
+----
+1  1
+
 statement error pgcode 0A000 FOR UPDATE is not allowed with set-returning functions in the target list
 SELECT generate_series(1, 2) FOR UPDATE
 
@@ -249,6 +307,15 @@ SELECT * FROM (SELECT generate_series(1, 2)) a 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, i 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, i FOR UPDATE OF a
+
+query II nosort
+SELECT * FROM (SELECT generate_series(1, 2)) a, i FOR UPDATE OF i
+----
+1  1
+2  1
+
 # Set-returning functions in the from list are allowed.
 query I nosort
 SELECT * FROM generate_series(1, 2) FOR UPDATE

pkg/sql/opt/optbuilder/builder.go

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

pkg/sql/opt/optbuilder/join.go

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