release-24.1: optbuilder: subject SELECT FOR UPDATE to sql_safe_updates

pkg/cli/interactive_tests/test_sql_safe_updates.tcl

@@ -19,7 +19,7 @@ eexpect "CREATE"
 eexpect root@
 
 send "delete from d.t;\r"
-eexpect "rejected (sql_safe_updates = true): DELETE without WHERE clause"
+eexpect "rejected (sql_safe_updates = true): DELETE without WHERE or LIMIT clause"
 eexpect root@
 end_test
 
@@ -35,7 +35,7 @@ send "show sql_safe_updates;\r"
 eexpect "on"
 eexpect "\r\n"
 send "delete from d.t;\r"
-eexpect "rejected (sql_safe_updates = true): DELETE without WHERE clause"
+eexpect "rejected (sql_safe_updates = true): DELETE without WHERE or LIMIT clause"
 send "\\q\r"
 eexpect ":/# "
 end_test
@@ -54,7 +54,7 @@ end_test
 
 start_test "Check that dangerous statements are properly rejected when using --safe-updates -e."
 send "$argv sql --safe-updates -e 'delete from d.t'\r"
-eexpect "rejected (sql_safe_updates = true): DELETE without WHERE clause"
+eexpect "rejected (sql_safe_updates = true): DELETE without WHERE or LIMIT clause"
 eexpect ":/# "
 end_test
 

pkg/sql/logictest/testdata/logic_test/dangerous_statements

@@ -4,12 +4,60 @@ CREATE TABLE foo(x INT)
 statement ok
 SET sql_safe_updates = true
 
-statement error rejected.*: UPDATE without WHERE clause
+statement error rejected.*: UPDATE without WHERE or LIMIT clause
 UPDATE foo SET x = 3
 
-statement error rejected.*: DELETE without WHERE clause
+statement ok
+UPDATE foo SET x = 3 WHERE x = 2
+
+statement ok
+UPDATE foo SET x = 3 ORDER BY x LIMIT 1
+
+statement error rejected.*: DELETE without WHERE or LIMIT clause
 DELETE FROM foo
 
+statement ok
+DELETE FROM foo WHERE x = 2
+
+statement ok
+DELETE FROM foo ORDER BY x LIMIT 1
+
+statement error rejected.*: SELECT FOR UPDATE without WHERE or LIMIT clause
+SELECT * FROM foo FOR UPDATE
+
+statement error rejected.*: SELECT FOR SHARE without WHERE or LIMIT clause
+SELECT * FROM foo FOR SHARE OF foo SKIP LOCKED
+
+statement ok
+SELECT * FROM foo WHERE x = 2 FOR UPDATE
+
+statement ok
+SELECT * FROM foo ORDER BY x LIMIT 1 FOR UPDATE
+
+statement error rejected.*: SELECT FOR UPDATE without WHERE or LIMIT clause
+(SELECT * FROM foo) FOR UPDATE
+
+statement ok
+(SELECT * FROM foo WHERE x = 2) FOR UPDATE
+
+statement ok
+SELECT * FROM (SELECT * FROM foo WHERE x = 2) FOR UPDATE
+
+statement ok
+SELECT * FROM (SELECT * FROM (SELECT * FROM foo) WHERE x = 2) FOR UPDATE
+
+statement error rejected.*: SELECT FOR UPDATE without WHERE or LIMIT clause
+SELECT * FROM (SELECT * FROM foo FOR UPDATE) WHERE x = 2 FOR UPDATE
+
+statement error rejected.*: SELECT FOR SHARE without WHERE or LIMIT clause
+SELECT * FROM (SELECT * FROM foo WHERE x = 2 FOR UPDATE) m, (SELECT * FROM foo) n FOR SHARE
+
+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
+
+statement ok
+SELECT * FROM (SELECT * FROM (SELECT * FROM foo) WHERE x > 1) WHERE x > 2 FOR UPDATE
+
 statement error rejected.*: ALTER TABLE DROP COLUMN
 ALTER TABLE foo DROP COLUMN x
 

pkg/sql/opt/optbuilder/delete.go

@@ -28,16 +28,16 @@ import (
 // mutations are applied, or the order of any returned rows (i.e. it won't
 // become a physical property required of the Delete operator).
 func (b *Builder) buildDelete(del *tree.Delete, inScope *scope) (outScope *scope) {
-	// UX friendliness safeguard.
-	if del.Where == nil && b.evalCtx.SessionData().SafeUpdates {
-		panic(pgerror.DangerousStatementf("DELETE without WHERE clause"))
-	}
-
 	if del.OrderBy != nil && del.Limit == nil {
 		panic(pgerror.Newf(pgcode.Syntax,
 			"DELETE statement requires LIMIT when ORDER BY is used"))
 	}
 
+	// UX friendliness safeguard.
+	if del.Where == nil && del.Limit == nil && b.evalCtx.SessionData().SafeUpdates {
+		panic(pgerror.DangerousStatementf("DELETE without WHERE or LIMIT clause"))
+	}
+
 	batch := del.Batch
 	if batch != nil {
 		var hasSize bool

pkg/sql/opt/optbuilder/locking.go

@@ -143,7 +143,9 @@ func (lm lockingSpec) get() opt.Locking {
 	return l
 }
 
-// lockingContext holds the locking information for the current scope.
+// lockingContext holds the locking information for the current scope. It is
+// passed down into subexpressions by value so that it automatically "pops" back
+// to its previous value on return.
 type lockingContext struct {
 	// lockScope is the stack of locking items that are currently in scope. This
 	// might include locking items that do not currently apply because they have
@@ -160,6 +162,12 @@ type lockingContext struct {
 	// return an error if the locking is set when we are building a table scan and
 	// isNullExtended is true.
 	isNullExtended bool
+
+	// safeUpdate is set to true if this lockingContext is being passed down from
+	// a select statement with either a WHERE clause or a LIMIT clause. This is
+	// needed so that we can return an error if we're locking without a WHERE
+	// clause or LIMIT clause and sql_safe_updates is true.
+	safeUpdate bool
 }
 
 // noLocking indicates that no row-level locking has been specified.

pkg/sql/opt/optbuilder/select.go

@@ -146,6 +146,13 @@ func (b *Builder) buildDataSource(
 					pgcode.FeatureNotSupported, "%s cannot be applied to the nullable side of an outer join",
 					lockCtx.locking.get().Strength))
 			}
+			// With sql_safe_updates set, we cannot use SELECT FOR UPDATE without a
+			// WHERE or LIMIT clause.
+			if !lockCtx.safeUpdate && b.evalCtx.SessionData().SafeUpdates {
+				panic(pgerror.DangerousStatementf(
+					"SELECT %s without WHERE or LIMIT clause", lockCtx.locking.get().Strength,
+				))
+			}
 			// SELECT ... FOR [KEY] UPDATE/SHARE also requires UPDATE privileges.
 			b.checkPrivilege(depName, ds, privilege.UPDATE)
 		}
@@ -275,6 +282,13 @@ func (b *Builder) buildDataSource(
 					pgcode.FeatureNotSupported, "%s cannot be applied to the nullable side of an outer join",
 					lockCtx.locking.get().Strength))
 			}
+			// With sql_safe_updates set, we cannot use SELECT FOR UPDATE without a
+			// WHERE or LIMIT clause.
+			if !lockCtx.safeUpdate && b.evalCtx.SessionData().SafeUpdates {
+				panic(pgerror.DangerousStatementf(
+					"SELECT %s without WHERE or LIMIT clause", lockCtx.locking.get().Strength,
+				))
+			}
 			// SELECT ... FOR [KEY] UPDATE/SHARE also requires UPDATE privileges.
 			b.checkPrivilege(depName, ds, privilege.UPDATE)
 		}
@@ -1136,6 +1150,12 @@ func (b *Builder) buildSelectStmtWithoutParens(
 	for i := len(lockingClause) - 1; i >= 0; i-- {
 		lockCtx.push(lockingClause[i])
 	}
+	if len(lockingClause) > 0 {
+		lockCtx.safeUpdate = false
+	}
+	if limit != nil {
+		lockCtx.safeUpdate = true
+	}
 
 	// NB: The case statements are sorted lexicographically.
 	switch t := wrapped.(type) {
@@ -1217,6 +1237,10 @@ func (b *Builder) buildSelectClause(
 	desiredTypes []*types.T,
 	inScope *scope,
 ) (outScope *scope) {
+	if sel.Where != nil {
+		lockCtx.safeUpdate = true
+	}
+
 	fromScope := b.buildFrom(sel.From, lockCtx, inScope)
 
 	b.processWindowDefs(sel, fromScope)

pkg/sql/opt/optbuilder/update.go

@@ -70,8 +70,8 @@ func (b *Builder) buildUpdate(upd *tree.Update, inScope *scope) (outScope *scope
 	}
 
 	// UX friendliness safeguard.
-	if upd.Where == nil && b.evalCtx.SessionData().SafeUpdates {
-		panic(pgerror.DangerousStatementf("UPDATE without WHERE clause"))
+	if upd.Where == nil && upd.Limit == nil && b.evalCtx.SessionData().SafeUpdates {
+		panic(pgerror.DangerousStatementf("UPDATE without WHERE or LIMIT clause"))
 	}
 
 	// Find which table we're working on, check the permissions.