move capture to rewrite

Signed-off-by: you06 <[email protected]>

planner/core/expression_rewriter.go

@@ -107,7 +107,11 @@ func (b *PlanBuilder) rewriteInsertOnDuplicateUpdate(ctx context.Context, exprNo
 // asScalar means whether this expression must be treated as a scalar expression.
 // And this function returns a result expression, a new plan that may have apply or semi-join.
 func (b *PlanBuilder) rewrite(ctx context.Context, exprNode ast.ExprNode, p LogicalPlan, aggMapper map[*ast.AggregateFuncExpr]int, asScalar bool) (expression.Expression, LogicalPlan, error) {
+	// capture the forUpdateRead mark
+	isForUpdateRead := b.isForUpdateRead
+	b.isForUpdateRead = false
 	expr, resultPlan, err := b.rewriteWithPreprocess(ctx, exprNode, p, aggMapper, nil, asScalar, nil)
+	b.isForUpdateRead = isForUpdateRead
 	return expr, resultPlan, err
 }
 

planner/core/logical_plan_builder.go

@@ -2235,11 +2235,7 @@ func (r *correlatedAggregateResolver) Enter(n ast.Node) (ast.Node, bool) {
 			r.b.outerSchemas = append(r.b.outerSchemas, outerSchema)
 			r.b.outerNames = append(r.b.outerNames, r.outerPlan.OutputNames())
 		}
-		// capture the forUpdateRead mark
-		isForUpdateRead := r.b.isForUpdateRead
-		r.b.isForUpdateRead = false
 		r.err = r.resolveSelect(v)
-		r.b.isForUpdateRead = isForUpdateRead
 		return n, true
 	}
 	return n, false

session/pessimistic_test.go

@@ -2400,8 +2400,13 @@ func (s *testPessimisticSuite) TestIssue21498(c *C) {
 		tk2.MustExec("update t set v = 33 where v = 32")
 		tk.MustExec("insert into t(id, v, v2) select 3 * id, 3 * v, 3 * v2 from t where v = 33")
 		tk.CheckExecResult(1, 0)
+		tk.MustExec("insert into t(id, v, v2) select (select 4 * id from t where v = 32) id, 4 * v, 4 * v2 from t where v = 33")
+		tk.CheckExecResult(1, 0)
+		err = tk.ExecToErr("insert into t(id, v, v2) select (select 4 * id from t where v = 33) id, 4 * v, 4 * v2 from t where v = 33")
+		c.Assert(err, NotNil)
+		c.Assert(err.Error(), Equals, "[table:1048]Column 'id' cannot be null")
 		tk.MustExec("commit")
-		tk.MustQuery("select * from t").Check(testkit.Rows("3 33 300", "5 11 100", "6 60 600", "9 99 900"))
+		tk.MustQuery("select * from t").Check(testkit.Rows("3 33 300", "5 11 100", "6 60 600", "9 99 900", "12 132 1200"))
 
 		tk2.MustExec("alter table t add unique index iv(v)")
 		tk2.MustExec("drop table if exists t1")