table: refactor options for `Table.AddRecord/Table.UpdateRecord/Index…

…Create (#55125)

close #55121

pkg/ddl/column_change_test.go

@@ -253,7 +253,7 @@ func checkAddWriteOnly(ctx sessionctx.Context, deleteOnlyTable, writeOnlyTable t
 	if err != nil {
 		return errors.Trace(err)
 	}
-	err = writeOnlyTable.UpdateRecord(context.Background(), ctx.GetTableCtx(), h, types.MakeDatums(1, 2, 3), types.MakeDatums(2, 2, 3), touchedSlice(writeOnlyTable))
+	err = writeOnlyTable.UpdateRecord(ctx.GetTableCtx(), h, types.MakeDatums(1, 2, 3), types.MakeDatums(2, 2, 3), touchedSlice(writeOnlyTable))
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -317,7 +317,7 @@ func checkAddPublic(sctx sessionctx.Context, writeOnlyTable, publicTable table.T
 		return errors.Errorf("%v", oldRow)
 	}
 	newRow := types.MakeDatums(3, 4, oldRow[2].GetValue())
-	err = writeOnlyTable.UpdateRecord(context.Background(), sctx.GetTableCtx(), h, oldRow, newRow, touchedSlice(writeOnlyTable))
+	err = writeOnlyTable.UpdateRecord(sctx.GetTableCtx(), h, oldRow, newRow, touchedSlice(writeOnlyTable))
 	if err != nil {
 		return errors.Trace(err)
 	}

pkg/ddl/index_change_test.go

@@ -173,7 +173,7 @@ func checkAddWriteOnlyForAddIndex(ctx sessionctx.Context, delOnlyTbl, writeOnlyT
 	}
 
 	// WriteOnlyTable: update t set c2 = 1 where c1 = 4 and c2 = 4
-	err = writeOnlyTbl.UpdateRecord(context.Background(), ctx.GetTableCtx(), kv.IntHandle(4), types.MakeDatums(4, 4), types.MakeDatums(4, 1), touchedSlice(writeOnlyTbl))
+	err = writeOnlyTbl.UpdateRecord(ctx.GetTableCtx(), kv.IntHandle(4), types.MakeDatums(4, 4), types.MakeDatums(4, 1), touchedSlice(writeOnlyTbl))
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -183,7 +183,7 @@ func checkAddWriteOnlyForAddIndex(ctx sessionctx.Context, delOnlyTbl, writeOnlyT
 	}
 
 	// DeleteOnlyTable: update t set c2 = 3 where c1 = 4 and c2 = 1
-	err = delOnlyTbl.UpdateRecord(context.Background(), ctx.GetTableCtx(), kv.IntHandle(4), types.MakeDatums(4, 1), types.MakeDatums(4, 3), touchedSlice(writeOnlyTbl))
+	err = delOnlyTbl.UpdateRecord(ctx.GetTableCtx(), kv.IntHandle(4), types.MakeDatums(4, 1), types.MakeDatums(4, 3), touchedSlice(writeOnlyTbl))
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -250,7 +250,7 @@ func checkAddPublicForAddIndex(ctx sessionctx.Context, writeTbl, publicTbl table
 	}
 
 	// WriteOnlyTable: update t set c2 = 5 where c1 = 7 and c2 = 7
-	err = writeTbl.UpdateRecord(context.Background(), ctx.GetTableCtx(), kv.IntHandle(7), types.MakeDatums(7, 7), types.MakeDatums(7, 5), touchedSlice(writeTbl))
+	err = writeTbl.UpdateRecord(ctx.GetTableCtx(), kv.IntHandle(7), types.MakeDatums(7, 7), types.MakeDatums(7, 5), touchedSlice(writeTbl))
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -328,7 +328,7 @@ func checkDropWriteOnly(ctx sessionctx.Context, publicTbl, writeTbl table.Table)
 	}
 
 	// WriteOnlyTable update t set c2 = 7 where c1 = 8 and c2 = 8
-	err = writeTbl.UpdateRecord(context.Background(), ctx.GetTableCtx(), kv.IntHandle(8), types.MakeDatums(8, 8), types.MakeDatums(8, 7), touchedSlice(writeTbl))
+	err = writeTbl.UpdateRecord(ctx.GetTableCtx(), kv.IntHandle(8), types.MakeDatums(8, 8), types.MakeDatums(8, 7), touchedSlice(writeTbl))
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -383,7 +383,7 @@ func checkDropDeleteOnly(ctx sessionctx.Context, writeTbl, delTbl table.Table) e
 	}
 
 	// DeleteOnlyTable update t set c2 = 10 where c1 = 9
-	err = delTbl.UpdateRecord(context.Background(), ctx.GetTableCtx(), kv.IntHandle(9), types.MakeDatums(9, 9), types.MakeDatums(9, 10), touchedSlice(delTbl))
+	err = delTbl.UpdateRecord(ctx.GetTableCtx(), kv.IntHandle(9), types.MakeDatums(9, 9), types.MakeDatums(9, 10), touchedSlice(delTbl))
 	if err != nil {
 		return errors.Trace(err)
 	}

pkg/executor/write.go

@@ -194,7 +194,7 @@ func updateRecord(
 		}
 	} else {
 		// Update record to new value and update index.
-		if err := t.UpdateRecord(ctx, sctx.GetTableCtx(), h, oldData, newData, modified); err != nil {
+		if err := t.UpdateRecord(sctx.GetTableCtx(), h, oldData, newData, modified, table.WithCtx(ctx)); err != nil {
 			if terr, ok := errors.Cause(err).(*terror.Error); ok && (terr.Code() == errno.ErrNoPartitionForGivenValue || terr.Code() == errno.ErrRowDoesNotMatchGivenPartitionSet) {
 				ec := sctx.GetSessionVars().StmtCtx.ErrCtx()
 				return false, ec.HandleError(err)

pkg/infoschema/tables.go

@@ -2436,7 +2436,7 @@ func (it *infoschemaTable) RemoveRecord(ctx table.MutateContext, h kv.Handle, r
 }
 
 // UpdateRecord implements table.Table UpdateRecord interface.
-func (it *infoschemaTable) UpdateRecord(gctx context.Context, ctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool) error {
+func (it *infoschemaTable) UpdateRecord(ctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool, opts ...table.UpdateRecordOption) error {
 	return table.ErrUnsupportedOp
 }
 
@@ -2529,7 +2529,7 @@ func (vt *VirtualTable) RemoveRecord(ctx table.MutateContext, h kv.Handle, r []t
 }
 
 // UpdateRecord implements table.Table UpdateRecord interface.
-func (vt *VirtualTable) UpdateRecord(ctx context.Context, sctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool) error {
+func (vt *VirtualTable) UpdateRecord(ctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool, opts ...table.UpdateRecordOption) error {
 	return table.ErrUnsupportedOp
 }
 

pkg/table/BUILD.bazel

@@ -53,7 +53,7 @@ go_test(
     embed = [":table"],
     flaky = True,
     race = "on",
-    shard_count = 9,
+    shard_count = 10,
     deps = [
         "//pkg/errctx",
         "//pkg/errno",

pkg/table/index.go

@@ -15,7 +15,6 @@
 package table
 
 import (
-	"context"
 	"time"
 
 	"github.com/pingcap/tidb/pkg/errctx"
@@ -32,51 +31,56 @@ type IndexIterator interface {
 
 // CreateIdxOpt contains the options will be used when creating an index.
 type CreateIdxOpt struct {
-	Ctx             context.Context
-	Untouched       bool // If true, the index key/value is no need to commit.
+	commonMutateOpt
 	IgnoreAssertion bool
 	FromBackFill    bool
 }
 
+// NewCreateIdxOpt creates a new CreateIdxOpt.
+func NewCreateIdxOpt(opts ...CreateIdxOption) *CreateIdxOpt {
+	opt := &CreateIdxOpt{}
+	for _, o := range opts {
+		o.ApplyCreateIdxOpt(opt)
+	}
+	return opt
+}
+
+// CreateIdxOption is defined for the Create() method of the Index interface.
+type CreateIdxOption interface {
+	ApplyCreateIdxOpt(*CreateIdxOpt)
+}
+
 // CreateIdxOptFunc is defined for the Create() method of Index interface.
 // Here is a blog post about how to use this pattern:
 // https://dave.cheney.net/2014/10/17/functional-options-for-friendly-apis
 type CreateIdxOptFunc func(*CreateIdxOpt)
 
-// IndexIsUntouched uses to indicate the index kv is untouched.
-var IndexIsUntouched CreateIdxOptFunc = func(opt *CreateIdxOpt) {
-	opt.Untouched = true
+// ApplyCreateIdxOpt implements the CreateIdxOption interface.
+func (f CreateIdxOptFunc) ApplyCreateIdxOpt(opt *CreateIdxOpt) {
+	f(opt)
 }
 
 // WithIgnoreAssertion uses to indicate the process can ignore assertion.
-var WithIgnoreAssertion = func(opt *CreateIdxOpt) {
+var WithIgnoreAssertion CreateIdxOptFunc = func(opt *CreateIdxOpt) {
 	opt.IgnoreAssertion = true
 }
 
 // FromBackfill indicates that the index is created by DDL backfill worker.
 // In the backfill-merge process, the index KVs from DML will be redirected to
 // the temp index. On the other hand, the index KVs from DDL backfill worker should
 // never be redirected to the temp index.
-var FromBackfill = func(opt *CreateIdxOpt) {
+var FromBackfill CreateIdxOptFunc = func(opt *CreateIdxOpt) {
 	opt.FromBackFill = true
 }
 
-// WithCtx returns a CreateIdxFunc.
-// This option is used to pass context.Context.
-func WithCtx(ctx context.Context) CreateIdxOptFunc {
-	return func(opt *CreateIdxOpt) {
-		opt.Ctx = ctx
-	}
-}
-
 // Index is the interface for index data on KV store.
 type Index interface {
 	// Meta returns IndexInfo.
 	Meta() *model.IndexInfo
 	// TableMeta returns TableInfo
 	TableMeta() *model.TableInfo
 	// Create supports insert into statement.
-	Create(ctx MutateContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle, handleRestoreData []types.Datum, opts ...CreateIdxOptFunc) (kv.Handle, error)
+	Create(ctx MutateContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle, handleRestoreData []types.Datum, opts ...CreateIdxOption) (kv.Handle, error)
 	// Delete supports delete from statement.
 	Delete(ctx MutateContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle) error
 	// GenIndexKVIter generate index key and value for multi-valued index, use iterator to reduce the memory allocation.

pkg/table/table.go

@@ -119,38 +119,106 @@ var (
 // RecordIterFunc is used for low-level record iteration.
 type RecordIterFunc func(h kv.Handle, rec []types.Datum, cols []*Column) (more bool, err error)
 
+// commonMutateOpt is the common options for mutating a table.
+type commonMutateOpt struct {
+	Ctx context.Context
+}
+
 // AddRecordOpt contains the options will be used when adding a record.
 type AddRecordOpt struct {
-	CreateIdxOpt
+	commonMutateOpt
 	IsUpdate      bool
 	ReserveAutoID int
 }
 
+// NewAddRecordOpt creates a new AddRecordOpt with options.
+func NewAddRecordOpt(opts ...AddRecordOption) *AddRecordOpt {
+	opt := &AddRecordOpt{}
+	for _, o := range opts {
+		o.ApplyAddRecordOpt(opt)
+	}
+	return opt
+}
+
+// GetCreateIdxOpt creates a CreateIdxOpt.
+func (opt *AddRecordOpt) GetCreateIdxOpt() *CreateIdxOpt {
+	return &CreateIdxOpt{commonMutateOpt: opt.commonMutateOpt}
+}
+
 // AddRecordOption is defined for the AddRecord() method of the Table interface.
 type AddRecordOption interface {
-	ApplyOn(*AddRecordOpt)
+	ApplyAddRecordOpt(*AddRecordOpt)
+}
+
+// UpdateRecordOpt contains the options will be used when updating a record.
+type UpdateRecordOpt struct {
+	commonMutateOpt
+}
+
+// NewUpdateRecordOpt creates a new UpdateRecordOpt with options.
+func NewUpdateRecordOpt(opts ...UpdateRecordOption) *UpdateRecordOpt {
+	opt := &UpdateRecordOpt{}
+	for _, o := range opts {
+		o.ApplyUpdateRecordOpt(opt)
+	}
+	return opt
+}
+
+// GetAddRecordOpt creates a AddRecordOpt.
+func (opt *UpdateRecordOpt) GetAddRecordOpt() *AddRecordOpt {
+	return &AddRecordOpt{commonMutateOpt: opt.commonMutateOpt}
+}
+
+// GetCreateIdxOpt creates a CreateIdxOpt.
+func (opt *UpdateRecordOpt) GetCreateIdxOpt() *CreateIdxOpt {
+	return &CreateIdxOpt{commonMutateOpt: opt.commonMutateOpt}
+}
+
+// UpdateRecordOption is defined for the UpdateRecord() method of the Table interface.
+type UpdateRecordOption interface {
+	ApplyUpdateRecordOpt(*UpdateRecordOpt)
+}
+
+// CommonMutateOptFunc is a function to provide common options for mutating a table.
+type CommonMutateOptFunc func(*commonMutateOpt)
+
+// ApplyAddRecordOpt implements the AddRecordOption interface.
+func (f CommonMutateOptFunc) ApplyAddRecordOpt(opt *AddRecordOpt) {
+	f(&opt.commonMutateOpt)
+}
+
+// ApplyUpdateRecordOpt implements the UpdateRecordOption interface.
+func (f CommonMutateOptFunc) ApplyUpdateRecordOpt(opt *UpdateRecordOpt) {
+	f(&opt.commonMutateOpt)
+}
+
+// ApplyCreateIdxOpt implements the CreateIdxOption interface.
+func (f CommonMutateOptFunc) ApplyCreateIdxOpt(opt *CreateIdxOpt) {
+	f(&opt.commonMutateOpt)
+}
+
+// WithCtx returns a CommonMutateOptFunc.
+// This option is used to pass context.Context.
+func WithCtx(ctx context.Context) CommonMutateOptFunc {
+	return func(opt *commonMutateOpt) {
+		opt.Ctx = ctx
+	}
 }
 
 // WithReserveAutoIDHint tells the AddRecord operation to reserve a batch of auto ID in the stmtctx.
 type WithReserveAutoIDHint int
 
-// ApplyOn implements the AddRecordOption interface.
-func (n WithReserveAutoIDHint) ApplyOn(opt *AddRecordOpt) {
+// ApplyAddRecordOpt implements the AddRecordOption interface.
+func (n WithReserveAutoIDHint) ApplyAddRecordOpt(opt *AddRecordOpt) {
 	opt.ReserveAutoID = int(n)
 }
 
-// ApplyOn implements the AddRecordOption interface, so any CreateIdxOptFunc
-// can be passed as the optional argument to the table.AddRecord method.
-func (f CreateIdxOptFunc) ApplyOn(opt *AddRecordOpt) {
-	f(&opt.CreateIdxOpt)
-}
-
 // IsUpdate is a defined value for AddRecordOptFunc.
 var IsUpdate AddRecordOption = isUpdate{}
 
 type isUpdate struct{}
 
-func (i isUpdate) ApplyOn(opt *AddRecordOpt) {
+func (i isUpdate) ApplyAddRecordOpt(opt *AddRecordOpt) {
 	opt.IsUpdate = true
 }
 
@@ -202,7 +270,7 @@ type Table interface {
 	AddRecord(ctx MutateContext, r []types.Datum, opts ...AddRecordOption) (recordID kv.Handle, err error)
 
 	// UpdateRecord updates a row which should contain only writable columns.
-	UpdateRecord(gctx context.Context, ctx MutateContext, h kv.Handle, currData, newData []types.Datum, touched []bool) error
+	UpdateRecord(ctx MutateContext, h kv.Handle, currData, newData []types.Datum, touched []bool, opts ...UpdateRecordOption) error
 
 	// RemoveRecord removes a row in the table.
 	RemoveRecord(ctx MutateContext, h kv.Handle, r []types.Datum) error

pkg/table/table_test.go

@@ -15,6 +15,7 @@
 package table
 
 import (
+	"context"
 	"testing"
 
 	mysql "github.com/pingcap/tidb/pkg/errno"
@@ -41,3 +42,39 @@ func TestErrorCode(t *testing.T) {
 	require.Equal(t, mysql.ErrNoPartitionForGivenValue, int(terror.ToSQLError(ErrNoPartitionForGivenValue).Code))
 	require.Equal(t, mysql.ErrLockOrActiveTransaction, int(terror.ToSQLError(ErrLockOrActiveTransaction).Code))
 }
+
+func TestOptions(t *testing.T) {
+	ctx := context.WithValue(context.Background(), "test", "test")
+	// NewAddRecordOpt without option
+	addOpt := NewAddRecordOpt()
+	require.Equal(t, AddRecordOpt{}, *addOpt)
+	require.Equal(t, CreateIdxOpt{}, *(addOpt.GetCreateIdxOpt()))
+	// NewAddRecordOpt with options
+	addOpt = NewAddRecordOpt(WithCtx(ctx), IsUpdate, WithReserveAutoIDHint(12))
+	require.Equal(t, AddRecordOpt{
+		commonMutateOpt: commonMutateOpt{Ctx: ctx},
+		IsUpdate:        true,
+		ReserveAutoID:   12,
+	}, *addOpt)
+	require.Equal(t, CreateIdxOpt{commonMutateOpt: commonMutateOpt{Ctx: ctx}}, *(addOpt.GetCreateIdxOpt()))
+	// NewUpdateRecordOpt without option
+	updateOpt := NewUpdateRecordOpt()
+	require.Equal(t, UpdateRecordOpt{}, *updateOpt)
+	require.Equal(t, AddRecordOpt{}, *(updateOpt.GetAddRecordOpt()))
+	require.Equal(t, CreateIdxOpt{}, *(updateOpt.GetCreateIdxOpt()))
+	// NewUpdateRecordOpt with options
+	updateOpt = NewUpdateRecordOpt(WithCtx(ctx))
+	require.Equal(t, UpdateRecordOpt{commonMutateOpt: commonMutateOpt{Ctx: ctx}}, *updateOpt)
+	require.Equal(t, AddRecordOpt{commonMutateOpt: commonMutateOpt{Ctx: ctx}}, *(updateOpt.GetAddRecordOpt()))
+	require.Equal(t, CreateIdxOpt{commonMutateOpt: commonMutateOpt{Ctx: ctx}}, *(updateOpt.GetCreateIdxOpt()))
+	// NewCreateIdxOpt without option
+	createIdxOpt := NewCreateIdxOpt()
+	require.Equal(t, CreateIdxOpt{}, *createIdxOpt)
+	// NewCreateIdxOpt with options
+	createIdxOpt = NewCreateIdxOpt(WithCtx(ctx), WithIgnoreAssertion, FromBackfill)
+	require.Equal(t, CreateIdxOpt{
+		commonMutateOpt: commonMutateOpt{Ctx: ctx},
+		IgnoreAssertion: true,
+		FromBackFill:    true,
+	}, *createIdxOpt)
+}

pkg/table/tables/BUILD.bazel

@@ -43,6 +43,7 @@ go_library(
         "//pkg/util/dbterror",
         "//pkg/util/generatedexpr",
         "//pkg/util/hack",
+        "//pkg/util/intest",
         "//pkg/util/logutil",
         "//pkg/util/ranger",
         "//pkg/util/rowcodec",

pkg/table/tables/bench_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/sessiontxn"
+	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/types"
 	_ "github.com/pingcap/tidb/pkg/util/context"
@@ -198,7 +199,7 @@ func BenchmarkUpdateRecordInPipelinedDML(b *testing.B) {
 		for j := 0; j < batchSize; j++ {
 			// Update record
 			handle := kv.IntHandle(j)
-			err := tb.UpdateRecord(context.TODO(), se.GetTableCtx(), handle, records[j], newData[j], touched)
+			err := tb.UpdateRecord(se.GetTableCtx(), handle, records[j], newData[j], touched, table.WithCtx(context.TODO()))
 			if err != nil {
 				b.Fatal(err)
 			}

pkg/table/tables/cache.go

@@ -254,13 +254,13 @@ func txnCtxAddCachedTable(sctx table.MutateContext, tid int64, handle *cachedTab
 }
 
 // UpdateRecord implements table.Table
-func (c *cachedTable) UpdateRecord(ctx context.Context, sctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool) error {
+func (c *cachedTable) UpdateRecord(ctx table.MutateContext, h kv.Handle, oldData, newData []types.Datum, touched []bool, opts ...table.UpdateRecordOption) error {
 	// Prevent furthur writing when the table is already too large.
 	if atomic.LoadInt64(&c.totalSize) > cachedTableSizeLimit {
 		return table.ErrOptOnCacheTable.GenWithStackByArgs("table too large")
 	}
-	txnCtxAddCachedTable(sctx, c.Meta().ID, c)
-	return c.TableCommon.UpdateRecord(ctx, sctx, h, oldData, newData, touched)
+	txnCtxAddCachedTable(ctx, c.Meta().ID, c)
+	return c.TableCommon.UpdateRecord(ctx, h, oldData, newData, touched, opts...)
 }
 
 // RemoveRecord implements table.Table RemoveRecord interface.