planner: fix update panic when update in prepare and execute

executor/executor.go

@@ -1685,6 +1685,17 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 			return err
 		}
 		s = prepareStmt.PreparedAst.Stmt
+		is := ctx.GetInfoSchema()
+		if prepareStmt.PreparedAst.SchemaVersion != is.SchemaMetaVersion() {
+			prepareStmt.PreparedAst.CachedPlan = nil
+			prepareStmt.Executor = nil
+			ret := &plannercore.PreprocessorReturn{InfoSchema: is.(infoschema.InfoSchema)}
+			err := plannercore.Preprocess(ctx, s, plannercore.InPrepare, plannercore.WithPreprocessorReturn(ret))
+			if err != nil {
+				return plannercore.ErrSchemaChanged.GenWithStack("Schema change caused error: %s", err.Error())
+			}
+			prepareStmt.PreparedAst.SchemaVersion = is.SchemaMetaVersion()
+		}
 		sc.InitSQLDigest(prepareStmt.NormalizedSQL, prepareStmt.SQLDigest)
 		// For `execute stmt` SQL, should reset the SQL digest with the prepare SQL digest.
 		goCtx := context.Background()

executor/prepared.go

@@ -329,6 +329,9 @@ func CompileExecutePreparedStmt(ctx context.Context, sctx sessionctx.Context,
 	if err != nil {
 		return nil, false, false, err
 	}
+	if is.SchemaMetaVersion() != sctx.GetInfoSchema().SchemaMetaVersion() {
+		is = sctx.GetInfoSchema().(infoschema.InfoSchema)
+	}
 
 	stmt := &ExecStmt{
 		GoCtx:       ctx,

planner/core/common_plans.go

@@ -282,6 +282,10 @@ func (e *Execute) OptimizePreparedPlan(ctx context.Context, sctx sessionctx.Cont
 		preparedObj.Executor = nil
 		// If the schema version has changed we need to preprocess it again,
 		// if this time it failed, the real reason for the error is schema changed.
+		// Example:
+		// When running update in prepared statement's schema version distinguished from the one of execute statement
+		// We should reset the tableRefs in the prepared update statements, otherwise, the ast nodes still hold the old
+		// tableRefs columnInfo which will cause chaos in logic of trying point get plan. (should ban non-public column)
 		ret := &PreprocessorReturn{InfoSchema: is}
 		err := Preprocess(sctx, prepared.Stmt, InPrepare, WithPreprocessorReturn(ret))
 		if err != nil {

planner/core/find_best_task.go

@@ -608,6 +608,20 @@ func (ds *DataSource) skylinePruning(prop *property.PhysicalProperty) []*candida
 	return candidates
 }
 
+func (ds *DataSource) isPointGetConvertableSchema() bool {
+	for _, col := range ds.Columns {
+		// Do not handle generated columns.
+		if col.IsGenerated() {
+			return false
+		}
+		// Only handle tables that all columns are public.
+		if col.State != model.StatePublic {
+			return false
+		}
+	}
+	return true
+}
+
 // findBestTask implements the PhysicalPlan interface.
 // It will enumerate all the available indices and choose a plan with least cost.
 func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter *PlanCounterTp) (t task, cntPlan int64, err error) {
@@ -704,7 +718,7 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 				p: dual,
 			}, cntPlan, nil
 		}
-		canConvertPointGet := len(path.Ranges) > 0 && path.StoreType == kv.TiKV
+		canConvertPointGet := len(path.Ranges) > 0 && path.StoreType == kv.TiKV && ds.isPointGetConvertableSchema()
 		if canConvertPointGet && !path.IsIntHandlePath {
 			// We simply do not build [batch] point get for prefix indexes. This can be optimized.
 			canConvertPointGet = path.Index.Unique && !path.Index.HasPrefixIndex()

planner/core/logical_plan_builder.go

@@ -4661,7 +4661,14 @@ func CheckUpdateList(assignFlags []int, updt *Update) error {
 		tbl := updt.tblID2Table[content.TblID]
 		flags := assignFlags[content.Start:content.End]
 		var update, updatePK bool
+		// step1: prepare s1 from "update t set a=? where b=?"   // use the schema before the alter table column
+		// step2: execute s1 using @a, @b                        // use the schema after the alter table column
+		// Since the prepare and execute use the difference schema, the step2 will see one more writable column
+		// --- changing column from ctc, which will lead the flag array out of range here.
 		for i, col := range tbl.WritableCols() {
+			//if col.ChangeStateInfo != nil {
+			//	continue
+			//}
 			if flags[i] >= 0 && col.State != model.StatePublic {
 				return ErrUnknownColumn.GenWithStackByArgs(col.Name, clauseMsg[fieldList])
 			}

planner/optimize.go

@@ -295,7 +295,21 @@ func optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 		}
 		if preparedPointer, ok := sctx.GetSessionVars().PreparedStmts[execID]; ok {
 			if preparedObj, ok := preparedPointer.(*core.CachedPrepareStmt); ok && preparedObj.ForUpdateRead {
+				// Actually, `IS` is read only for all the travel path around optimize and exec. The `IS` and `TxnCtn.IS`
+				// should always be the same, otherwise, it will cause some check fails. For example. `checkUpdateList`
+				// panics in the plan building phase.
+				//
+				//  ...
+				//  optimize(is) -> OptimizePreparedPlan(is) -> ... `is` changed
+				//  runstmt(is)
+				//  ...
+				// The modification of `is` is in the deep call chain. once func the backtracking to the outer, since
+				// the `is` struct value copy model, the outer `is` is still the old one before calling the runstmt.
+				//
+				// while the session TxnCtx hold the new `IS`, we can substitute the outer `is` once we found the schema
+				// version difference before calling the runstmt
 				is = domain.GetDomain(sctx).InfoSchema()
+				sctx.GetSessionVars().TxnCtx.InfoSchema = is
 			}
 		}
 	}

session/session.go

@@ -227,6 +227,10 @@ type session struct {
 	cache [1]ast.StmtNode
 
 	builtinFunctionUsage telemetry.BuiltinFunctionsUsage
+
+	debug map[string]string
+
+	cnt int
 }
 
 var parserPool = &sync.Pool{New: func() interface{} { return parser.New() }}
@@ -2047,6 +2051,8 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 	} else {
 		is = s.GetInfoSchema().(infoschema.InfoSchema)
 	}
+	// sync the new information schema with the one in TxnCtx of sessionVar.
+	s.sessionVars.TxnCtx.InfoSchema = is
 	if ok {
 		return s.cachedPlanExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
 	}