plan: get correct columns (#5846)

ddl/ddl_db_change_test.go

@@ -285,6 +285,64 @@ func (t *testExecInfo) execSQL(idx int) error {
 	return nil
 }
 
+// TestUpdateOrDelete tests whether the correct columns is used in PhysicalIndexScan's ToPB function.
+func (s *testStateChangeSuite) TestUpdateOrDelete(c *C) {
+	sqls := make([]string, 2)
+	sqls[0] = "delete from t where c2 = 'a'"
+	sqls[1] = "update t use index(c2) set c2 = 'c2_update' where c2 = 'a'"
+	alterTableSQL := "alter table t add column a int not null default 1 first"
+	s.runTestInWriteOnly(c, "", alterTableSQL, sqls)
+}
+
+func (s *testStateChangeSuite) runTestInWriteOnly(c *C, tableName, alterTableSQL string, sqls []string) {
+	defer testleak.AfterTest(c)()
+	_, err := s.se.Execute(goctx.Background(), `create table t (
+		c1 int primary key,
+		c2 varchar(64),
+		c3 enum('N','Y') not null default 'N',
+		c4 timestamp on update current_timestamp,
+		unique key(c2))`)
+	c.Assert(err, IsNil)
+	defer s.se.Execute(goctx.Background(), "drop table t")
+	_, err = s.se.Execute(goctx.Background(), "insert into t values(8, 'a', 'N', '2017-07-01')")
+	c.Assert(err, IsNil)
+	// Make sure these sqls use the the plan of index scan.
+	_, err = s.se.Execute(goctx.Background(), "drop stats t")
+	c.Assert(err, IsNil)
+
+	callback := &ddl.TestDDLCallback{}
+	prevState := model.StateNone
+	var checkErr error
+	times := 0
+	se, err := tidb.CreateSession(s.store)
+	c.Assert(err, IsNil)
+	_, err = se.Execute(goctx.Background(), "use test_db_state")
+	c.Assert(err, IsNil)
+	callback.OnJobUpdatedExported = func(job *model.Job) {
+		if job.SchemaState == prevState || checkErr != nil || times >= 3 {
+			return
+		}
+		times++
+		if job.SchemaState != model.StateWriteOnly {
+			return
+		}
+		for _, sql := range sqls {
+			_, err = se.Execute(goctx.Background(), sql)
+			if err != nil {
+				checkErr = err
+				break
+			}
+		}
+	}
+	d := s.dom.DDL()
+	d.SetHook(callback)
+	_, err = s.se.Execute(goctx.Background(), alterTableSQL)
+	c.Assert(err, IsNil)
+	c.Assert(errors.ErrorStack(checkErr), Equals, "")
+	callback = &ddl.TestDDLCallback{}
+	d.SetHook(callback)
+}
+
 func (s *testStateChangeSuite) execQuery(tk *testkit.TestKit, sql string, args ...interface{}) (*testkit.Result, error) {
 	comment := Commentf("sql:%s, args:%v", sql, args)
 	rs, err := tk.Exec(sql, args...)

model/model.go

@@ -181,6 +181,22 @@ func (t *TableInfo) GetPkColInfo() *ColumnInfo {
 	return nil
 }
 
+// Cols returns the columns of the table in public state.
+func (t *TableInfo) Cols() []*ColumnInfo {
+	publicColumns := make([]*ColumnInfo, len(t.Columns))
+	maxOffset := -1
+	for _, col := range t.Columns {
+		if col.State != StatePublic {
+			continue
+		}
+		publicColumns[col.Offset] = col
+		if maxOffset < col.Offset {
+			maxOffset = col.Offset
+		}
+	}
+	return publicColumns[0 : maxOffset+1]
+}
+
 // NewExtraHandleColInfo mocks a column info for extra handle column.
 func NewExtraHandleColInfo() *ColumnInfo {
 	colInfo := &ColumnInfo{

plan/plan_to_pb.go

@@ -118,14 +118,15 @@ func checkCoverIndex(idx *model.IndexInfo, ranges []*ranger.NewRange) bool {
 // ToPB implements PhysicalPlan ToPB interface.
 func (p *PhysicalIndexScan) ToPB(ctx context.Context) (*tipb.Executor, error) {
 	columns := make([]*model.ColumnInfo, 0, p.schema.Len())
+	tableColumns := p.Table.Cols()
 	for _, col := range p.schema.Columns {
 		if col.ID == model.ExtraHandleID {
 			columns = append(columns, &model.ColumnInfo{
 				ID:   model.ExtraHandleID,
 				Name: model.NewCIStr("_rowid"),
 			})
 		} else {
-			columns = append(columns, p.Table.Columns[col.Position])
+			columns = append(columns, tableColumns[col.Position])
 		}
 	}
 	idxExec := &tipb.IndexScan{