optbuilder: refactor semantic analysis of FOR UPDATE

pkg/sql/logictest/testdata/logic_test/select_for_update

@@ -131,6 +131,29 @@ query I
 ----
 1
 
+# Unlike Postgres, we do not support multiple locking clauses on parenthesized
+# queries.
+
+statement error pgcode 42601 multiple FOR UPDATE clauses not allowed
+(SELECT 1 FOR UPDATE) FOR UPDATE
+
+statement error pgcode 42601 multiple FOR UPDATE clauses not allowed
+((SELECT 1 FOR UPDATE)) FOR UPDATE
+
+statement error pgcode 42601 multiple FOR UPDATE clauses not allowed
+((SELECT 1) FOR UPDATE) FOR UPDATE
+
+# But we do support locking clauses both inside and outside subqueries.
+
+statement ok
+SELECT (SELECT 1 FOR UPDATE) FOR UPDATE
+
+statement ok
+SELECT * FROM (SELECT 1 FOR UPDATE) AS x FOR UPDATE
+
+statement ok
+SELECT * FROM (SELECT 1 FOR UPDATE) AS x WHERE EXISTS (SELECT 1 FOR UPDATE) FOR UPDATE
+
 # FOR READ ONLY is ignored, like in Postgres.
 query I
 SELECT 1 FOR READ ONLY
@@ -140,54 +163,159 @@ SELECT 1 FOR READ ONLY
 # Various operations are not supported when locking clauses are provided.
 # FeatureNotSupported errors are thrown for each of them.
 
+statement ok
+CREATE TABLE i (i PRIMARY KEY) AS SELECT 1
+
 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 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
 
 statement error pgcode 0A000 FOR UPDATE is not allowed with VALUES
+(VALUES (1)) FOR UPDATE
+
+# VALUES within a subquery is allowed, like in Postgres.
+
+query I
+SELECT (VALUES (1)) FOR UPDATE
+----
+1
+
+query I
 SELECT * FROM (VALUES (1)) a FOR UPDATE
+----
+1
+
+query II
+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
 
 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 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
 
 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 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
 
 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 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
 
 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 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
 
 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 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
 
 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
 
+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
@@ -201,6 +329,12 @@ SELECT * FROM (SELECT * FROM generate_series(1, 2)) a FOR UPDATE
 1
 2
 
+query II nosort
+SELECT * FROM (SELECT * FROM generate_series(1, 2)) a, i FOR UPDATE
+----
+1  1
+2  1
+
 # Use of SELECT FOR UPDATE/SHARE requires SELECT and UPDATE privileges.
 
 statement ok

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)