plan: get correct columns

ddl/ddl_db_change_test.go

@@ -48,6 +48,7 @@ func (s *testStateChangeSuite) SetUpSuite(c *C) {
 	s.store, err = tikv.NewMockTikvStore()
 	c.Assert(err, IsNil)
 	tidb.SetSchemaLease(s.lease)
+	tidb.SetStatsLease(0)
 	s.dom, err = tidb.BootstrapSession(s.store)
 	c.Assert(err, IsNil)
 	s.se, err = tidb.CreateSession(s.store)
@@ -281,3 +282,61 @@ 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 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(`create table t (
+		c1 int primary key,
+		c2 varchar(64),
+		c3 enum('N','Y') not null default 'N',
+		c4 timestamp on update current_timestamp,
+		key(c2))`)
+	c.Assert(err, IsNil)
+	defer s.se.Execute("drop table t")
+	_, err = s.se.Execute("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("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("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(sql)
+			if err != nil {
+				checkErr = err
+				break
+			}
+		}
+	}
+	d := s.dom.DDL()
+	d.SetHook(callback)
+	_, err = s.se.Execute(alterTableSQL)
+	c.Assert(err, IsNil)
+	c.Assert(errors.ErrorStack(checkErr), Equals, "")
+	callback = &ddl.TestDDLCallback{}
+	d.SetHook(callback)
+}

executor/prepared.go

@@ -360,6 +360,7 @@ func ResetStmtCtx(ctx context.Context, s ast.StmtNode) {
 		sc.IgnoreTruncate = false
 		sc.OverflowAsWarning = false
 		sc.TruncateAsWarning = !sessVars.StrictSQLMode || stmt.IgnoreErr
+		sc.InUpdateStmt = true
 		sc.InUpdateOrDeleteStmt = true
 		sc.DividedByZeroAsWarning = stmt.IgnoreErr
 		sc.IgnoreZeroInDate = !sessVars.StrictSQLMode || stmt.IgnoreErr

model/model.go

@@ -172,6 +172,39 @@ 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]
+}
+
+// WritableCols returns columns of the table in writable states.
+// Writable states includes Public, WriteOnly, WriteOnlyReorganization.
+func (t *TableInfo) WritableCols() []*ColumnInfo {
+	writableColumns := make([]*ColumnInfo, len(t.Columns))
+	maxOffset := -1
+	for _, col := range t.Columns {
+		if col.State == StateDeleteOnly || col.State == StateDeleteReorganization {
+			continue
+		}
+		writableColumns[col.Offset] = col
+		if maxOffset < col.Offset {
+			maxOffset = col.Offset
+		}
+	}
+	return writableColumns[0 : maxOffset+1]
+}
+
 // NewExtraHandleColInfo mocks a column info for extra handle column.
 func NewExtraHandleColInfo() *ColumnInfo {
 	colInfo := &ColumnInfo{

plan/logical_plan_builder.go

@@ -1251,7 +1251,8 @@ func (b *planBuilder) buildDataSource(tn *ast.TableName) LogicalPlan {
 	b.visitInfo = appendVisitInfo(b.visitInfo, mysql.SelectPriv, schemaName.L, tableInfo.Name.L, "")
 
 	var columns []*table.Column
-	if b.inUpdateStmt {
+	sc := b.ctx.GetSessionVars().StmtCtx
+	if sc.InUpdateStmt {
 		columns = tbl.WritableCols()
 	} else {
 		columns = tbl.Cols()
@@ -1454,7 +1455,6 @@ func (b *planBuilder) buildSemiJoin(outerPlan, innerPlan LogicalPlan, onConditio
 }
 
 func (b *planBuilder) buildUpdate(update *ast.UpdateStmt) LogicalPlan {
-	b.inUpdateStmt = true
 	b.needColHandle++
 	sel := &ast.SelectStmt{Fields: &ast.FieldList{}, From: update.TableRefs, Where: update.Where, OrderBy: update.Order, Limit: update.Limit}
 	p := b.buildResultSetNode(sel.From.TableRefs)

plan/plan_to_pb.go

@@ -89,14 +89,18 @@ func (p *PhysicalTableScan) ToPB(ctx context.Context) (*tipb.Executor, error) {
 // 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()
+	if ctx.GetSessionVars().StmtCtx.InUpdateStmt {
+		tableColumns = p.Table.WritableCols()
+	}
 	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{

plan/planbuilder.go

@@ -123,7 +123,6 @@ type planBuilder struct {
 	ctx           context.Context
 	is            infoschema.InfoSchema
 	outerSchemas  []*expression.Schema
-	inUpdateStmt  bool
 	needColHandle int
 	// colMapper stores the column that must be pre-resolved.
 	colMapper map[*ast.ColumnNameExpr]int

sessionctx/variable/session.go

@@ -339,6 +339,7 @@ type StatementContext struct {
 	// Set the following variables before execution
 
 	InInsertStmt           bool
+	InUpdateStmt           bool
 	InUpdateOrDeleteStmt   bool
 	InSelectStmt           bool
 	IgnoreTruncate         bool