diff --git a/lightning/pkg/errormanager/errormanager.go b/lightning/pkg/errormanager/errormanager.go index 3ef21f9fb0435..ef3ddf4e4cd6f 100644 --- a/lightning/pkg/errormanager/errormanager.go +++ b/lightning/pkg/errormanager/errormanager.go @@ -616,7 +616,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), overwrittenHandle, tbl.Cols(), overwritten) + tbl, overwrittenHandle, tbl.Cols(), overwritten) if err != nil { return errors.Trace(err) } @@ -796,7 +796,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), handle, tbl.Cols(), latestValue) + tbl, handle, tbl.Cols(), latestValue) if err != nil { return errors.Trace(err) } @@ -825,7 +825,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), handle, tbl.Cols(), rawValue) + tbl, handle, tbl.Cols(), rawValue) if err != nil { return errors.Trace(err) } diff --git a/pkg/ddl/backfilling_dist_executor.go b/pkg/ddl/backfilling_dist_executor.go index 0403d4ad0e854..763299c3ea0b8 100644 --- a/pkg/ddl/backfilling_dist_executor.go +++ b/pkg/ddl/backfilling_dist_executor.go @@ -19,6 +19,7 @@ import ( "encoding/json" "github.com/pingcap/errors" + "github.com/pingcap/failpoint" "github.com/pingcap/tidb/pkg/ddl/logutil" sess "github.com/pingcap/tidb/pkg/ddl/session" "github.com/pingcap/tidb/pkg/dxf/framework/proto" @@ -137,9 +138,10 @@ func (s *backfillDistExecutor) newBackfillStepExecutor( store := s.TaskRuntime.Store() sessPool := sess.NewSessionPool(s.TaskRuntime.SysSessionPool()) - // TODO getTableByTxn is using DDL ctx which is never cancelled except when shutdown. + // TODO This is using DDL ctx which is never cancelled except when shutdown. // we should move this operation out of GetStepExecutor, and put into Init. - _, tblIface, err := getTableByTxn(ddlObj.ctx, store, jobMeta.SchemaID, jobMeta.TableID) + failpoint.InjectCall("beforeGetUserTableForBackfillStep", jobMeta) + tblIface, err := getUserTableFromTaskStore(ddlObj.ctx, store, jobMeta) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_dist_scheduler.go b/pkg/ddl/backfilling_dist_scheduler.go index 46a1a98d34505..7016592b2b268 100644 --- a/pkg/ddl/backfilling_dist_scheduler.go +++ b/pkg/ddl/backfilling_dist_scheduler.go @@ -47,7 +47,9 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/store/helper" "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/util/backoff" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/tikv/client-go/v2/oracle" "github.com/tikv/client-go/v2/tikv" "go.uber.org/zap" @@ -200,7 +202,17 @@ func getUserTableFromTaskStore( return nil, err } // we don't touch table data during add-index, a fake Allocators is enough. - tbl, err := table.TableFromMeta(autoid.NewAllocators(tblInfo.SepAutoInc()), tblInfo) + defaultUseNewCollate := collate.NewCollationEnabled() + failpoint.Inject("overrideDefaultUseNewCollateForBackfillStep", func(val failpoint.Value) { + defaultUseNewCollate = val.(bool) + }) + useNewCollate := job.ReorgMeta.GetUseNewCollateOrDefault(defaultUseNewCollate) + failpoint.InjectCall("afterResolveUserTableNewCollateForBackfillStep", job, defaultUseNewCollate, useNewCollate) + tbl, err := tables.TableFromMetaWithCollate( + useNewCollate, + autoid.NewAllocators(tblInfo.SepAutoInc()), + tblInfo, + ) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index 4af509e6baf51..80d8ac9f285e8 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/pkg/ddl/backfilling_operators.go @@ -111,7 +111,7 @@ func NewAddIndexIngestPipeline( ) (*operator.AsyncPipeline, error) { indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndex(tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -169,7 +169,7 @@ func NewWriteIndexToExternalStoragePipeline( ) (*operator.AsyncPipeline, error) { indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndex(tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -914,7 +914,8 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i // skip running the checker in TiDB side. indexConditionCheckers = nil } - cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta()) + cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, + sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), w.tbl.UseNewCollate()) if err != nil || cnt == 0 { return 0, 0, err } diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 46c2121ce6acf..f2a72e3fc1382 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/intest" @@ -142,6 +143,7 @@ func NewReorgCopContext( tblInfo, allIdxInfo, requestSource, + reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), ) } diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index d39ac039a1b94..26cab165d69e4 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -46,6 +46,7 @@ type CopContextBase struct { ExprCtx exprctx.BuildContext PushDownFlags uint64 RequestSource string + UseNewCollate bool ColumnInfos []*model.ColumnInfo FieldTypes []*types.FieldType @@ -80,6 +81,7 @@ func NewCopContextBase( tblInfo *model.TableInfo, idxCols []*model.IndexColumn, requestSource string, + useNewCollate bool, ) (*CopContextBase, error) { var err error usedColumnIDs := make(map[int64]struct{}, len(idxCols)) @@ -125,8 +127,14 @@ func NewCopContextBase( handleIDs = []int64{extra.ID} } - expColInfos, _, err := expression.ColumnInfos2ColumnsAndNames(exprCtx, - ast.CIStr{} /* unused */, tblInfo.Name, colInfos, tblInfo) + expColInfos, _, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( + exprCtx, + ast.CIStr{}, // unused + tblInfo.Name, + colInfos, + tblInfo, + useNewCollate, + ) if err != nil { return nil, err } @@ -139,6 +147,7 @@ func NewCopContextBase( ExprCtx: exprCtx, PushDownFlags: pushDownFlags, RequestSource: requestSource, + UseNewCollate: useNewCollate, ColumnInfos: colInfos, FieldTypes: fieldTps, ExprColumnInfos: expColInfos, @@ -148,27 +157,36 @@ func NewCopContextBase( }, nil } -// NewCopContext creates a CopContext. +// NewCopContext creates a CopContext with a fixed collation mode. func NewCopContext( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, + useNewCollate bool, ) (CopContext, error) { if len(allIdxInfo) == 1 { - return NewCopContextSingleIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo[0], requestSource) + return NewCopContextSingleIndex( + exprCtx, + pushDownFlags, + tblInfo, + allIdxInfo[0], + requestSource, + useNewCollate, + ) } - return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource) + return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, useNewCollate) } -// NewCopContextSingleIndex creates a CopContextSingleIndex. +// NewCopContextSingleIndex creates a CopContextSingleIndex with a fixed collation mode. func NewCopContextSingleIndex( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, requestSource string, + useNewCollate bool, ) (*CopContextSingleIndex, error) { cols := idxInfo.Columns neededCols, err := tables.ExtractColumnsFromCondition(exprCtx, idxInfo, tblInfo, false) @@ -178,7 +196,7 @@ func NewCopContextSingleIndex( cols = append(cols, neededCols...) cols = tables.DedupIndexColumns(cols) - base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource) + base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource, useNewCollate) if err != nil { return nil, err } @@ -228,13 +246,14 @@ func (c *CopContextSingleIndex) GetCondition() (expression.Expression, error) { return expr, nil } -// NewCopContextMultiIndex creates a CopContextMultiIndex. +// NewCopContextMultiIndex creates a CopContextMultiIndex with a fixed collation mode. func NewCopContextMultiIndex( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, + useNewCollate bool, ) (*CopContextMultiIndex, error) { approxColLen := 0 for _, idxInfo := range allIdxInfo { @@ -252,7 +271,7 @@ func NewCopContextMultiIndex( } allIdxCols = tables.DedupIndexColumns(allIdxCols) - base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource) + base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource, useNewCollate) if err != nil { return nil, err } diff --git a/pkg/ddl/copr/copr_ctx_test.go b/pkg/ddl/copr/copr_ctx_test.go index 0e9a9caa0e8a2..80b4a676f63c9 100644 --- a/pkg/ddl/copr/copr_ctx_test.go +++ b/pkg/ddl/copr/copr_ctx_test.go @@ -110,7 +110,10 @@ func TestNewCopContextSingleIndex(t *testing.T) { copCtx, err := NewCopContextSingleIndex( sctx.GetExprCtx(), sctx.GetSessionVars().StmtCtx.PushDownFlags(), - mockTableInfo, mockIdxInfo, "", + mockTableInfo, + mockIdxInfo, + "", + false, ) require.NoError(t, err) base := copCtx.GetBase() diff --git a/pkg/ddl/executor.go b/pkg/ddl/executor.go index ea8f170051857..804be37b7bcca 100644 --- a/pkg/ddl/executor.go +++ b/pkg/ddl/executor.go @@ -7535,6 +7535,7 @@ func getJobCheckInterval(action model.ActionType, i int) (time.Duration, bool) { // NewDDLReorgMeta create a DDL ReorgMeta. func NewDDLReorgMeta(ctx sessionctx.Context) *model.DDLReorgMeta { tzName, tzOffset := ddlutil.GetTimeZone(ctx) + useNewCollate := collate.NewCollationEnabled() return &model.DDLReorgMeta{ SQLMode: ctx.GetSessionVars().SQLMode, Warnings: make(map[errors.ErrorID]*terror.Error), @@ -7542,6 +7543,7 @@ func NewDDLReorgMeta(ctx sessionctx.Context) *model.DDLReorgMeta { Location: &model.TimeZoneLocation{Name: tzName, Offset: tzOffset}, ResourceGroupName: ctx.GetSessionVars().StmtCtx.ResourceGroupName, Version: model.CurrentReorgMetaVersion, + UseNewCollate: &useNewCollate, } } diff --git a/pkg/ddl/export_test.go b/pkg/ddl/export_test.go index 40d13e07a14e7..2b55667c3e4f7 100644 --- a/pkg/ddl/export_test.go +++ b/pkg/ddl/export_test.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/mock" ) @@ -80,6 +81,6 @@ func ConvertRowToHandleAndIndexDatum( c := copCtx.GetBase() idxData := ddl.ExtractDatumByOffsets(ctx, row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, idxDataBuf) handleData := ddl.ExtractDatumByOffsets(ctx, row, c.HandleOutputOffsets, c.ExprColumnInfos, handleDataBuf) - handle, err := ddl.BuildHandle(handleData, c.TableInfo, c.PrimaryKeyInfo, time.Local, errctx.StrictNoWarningContext) + handle, err := ddl.BuildHandle(collate.NewCollationEnabled(), handleData, c.TableInfo, c.PrimaryKeyInfo, time.Local, errctx.StrictNoWarningContext) return handle, idxData, err } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 863fabffbdde8..f675c705921dd 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -82,7 +82,6 @@ import ( "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/backoff" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/engine" "github.com/pingcap/tidb/pkg/util/generatedexpr" @@ -2479,7 +2478,7 @@ func (w *baseIndexWorker) getIndexRecord(idxInfo *model.IndexInfo, handle kv.Han idxVal[j] = idxColumnVal } - rsData := tables.TryGetHandleRestoredDataWrapper(collate.NewCollationEnabled(), w.table.Meta(), nil, w.rowMap, idxInfo) + rsData := tables.TryGetHandleRestoredDataWrapper(w.table, nil, w.rowMap, idxInfo) idxRecord := &indexRecord{handle: handle, key: recordKey, vals: idxVal, rsData: rsData} return idxRecord, nil } @@ -2709,6 +2708,7 @@ func writeChunk( writeStmtBufs *variable.WriteStmtBufs, copChunk *chunk.Chunk, tblInfo *model.TableInfo, + useNewCollate bool, ) (rowCnt int, bytes int, err error) { iter := chunk.NewIterator4Chunk(copChunk) c := copCtx.GetBase() @@ -2733,7 +2733,6 @@ func writeChunk( }() needRestoreForIndexes := make([]bool, len(indexes)) restore, pkNeedRestore := false, false - useNewCollate := collate.NewCollationEnabled() if c.PrimaryKeyInfo != nil && c.TableInfo.IsCommonHandle && c.TableInfo.CommonHandleVersion != 0 { pkNeedRestore = tables.NeedRestoredData(useNewCollate, c.PrimaryKeyInfo.Columns, c.TableInfo.Columns) } @@ -2754,7 +2753,7 @@ func writeChunk( restoreDataBuf[i] = *datum.Clone() } } - h, err := BuildHandle(handleDataBuf, c.TableInfo, c.PrimaryKeyInfo, loc, errCtx) + h, err := BuildHandle(useNewCollate, handleDataBuf, c.TableInfo, c.PrimaryKeyInfo, loc, errCtx) if err != nil { return 0, totalBytes, errors.Trace(err) } @@ -2777,7 +2776,7 @@ func writeChunk( idxData := idxDataBuf[:len(index.Meta().Columns)] var rsData []types.Datum if needRestoreForIndexes[i] { - rsData = getRestoreData(c.TableInfo, copCtx.IndexInfo(idxID), c.PrimaryKeyInfo, restoreDataBuf) + rsData = getRestoreData(useNewCollate, c.TableInfo, copCtx.IndexInfo(idxID), c.PrimaryKeyInfo, restoreDataBuf) } kvBytes, err := writeOneKV(ctx, writers[i], index, loc, errCtx, writeStmtBufs, idxData, rsData, h) if err != nil { diff --git a/pkg/ddl/index_cop.go b/pkg/ddl/index_cop.go index baafbf4f5149d..65adf1ee8e7c6 100644 --- a/pkg/ddl/index_cop.go +++ b/pkg/ddl/index_cop.go @@ -36,7 +36,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/timeutil" "github.com/pingcap/tipb/go-tipb" @@ -146,8 +145,8 @@ func completeErr(err error, idxInfo *model.IndexInfo) error { return errors.Trace(err) } -func getRestoreData(tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, handleDts []types.Datum) []types.Datum { - if !collate.NewCollationEnabled() || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { +func getRestoreData(useNewCollate bool, tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, handleDts []types.Datum) []types.Datum { + if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { return nil } if pkIdx == nil { @@ -155,7 +154,7 @@ func getRestoreData(tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, } for i, pkIdxCol := range pkIdx.Columns { pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredData(&pkCol.FieldType) { + if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { // Since the handle data cannot be null, we can use SetNull to // indicate that this column does not need to be restored. handleDts[i].SetNull() @@ -257,11 +256,11 @@ func ExtractDatumByOffsets(ctx expression.EvalContext, row chunk.Row, offsets [] } // BuildHandle is exported for test. -func BuildHandle(pkDts []types.Datum, tblInfo *model.TableInfo, +func BuildHandle(useNewCollate bool, pkDts []types.Datum, tblInfo *model.TableInfo, pkInfo *model.IndexInfo, loc *time.Location, errCtx errctx.Context) (kv.Handle, error) { if tblInfo.IsCommonHandle { tablecodec.TruncateIndexValues(tblInfo, pkInfo, pkDts) - handleBytes, err := codec.EncodeKey(loc, nil, pkDts...) + handleBytes, err := codec.NewEncoder(useNewCollate).EncodeKey(loc, nil, pkDts...) err = errCtx.HandleError(err) if err != nil { return nil, err diff --git a/pkg/ddl/job_scheduler.go b/pkg/ddl/job_scheduler.go index 7cbbd6f8a7913..752e48c92ceed 100644 --- a/pkg/ddl/job_scheduler.go +++ b/pkg/ddl/job_scheduler.go @@ -40,13 +40,11 @@ import ( "github.com/pingcap/tidb/pkg/ddl/util" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta" - "github.com/pingcap/tidb/pkg/meta/autoid" "github.com/pingcap/tidb/pkg/meta/metadef" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/metrics" "github.com/pingcap/tidb/pkg/owner" "github.com/pingcap/tidb/pkg/sessionctx/vardef" - "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/etcd" @@ -691,29 +689,6 @@ func (s *jobScheduler) workerPoolExhausted() bool { s.bgJobWorkerPool.available() == 0 } -func getTableByTxn(ctx context.Context, store kv.Storage, schemaID, tableID int64) (*model.DBInfo, table.Table, error) { - var tbl table.Table - var dbInfo *model.DBInfo - err := kv.RunInNewTxn(ctx, store, false, func(_ context.Context, txn kv.Transaction) error { - t := meta.NewMutator(txn) - var err1 error - dbInfo, err1 = t.GetDatabase(schemaID) - if err1 != nil { - return errors.Trace(err1) - } - tblInfo, err1 := getTableInfo(t, tableID, schemaID) - if err1 != nil { - return errors.Trace(err1) - } - // This tableInfo should never interact with the autoid allocator, - // so we can use the autoid.Allocators{} here. - // TODO(tangenta): Use model.TableInfo instead of tables.Table. - tbl, err1 = table.TableFromMeta(autoid.Allocators{}, tblInfo) - return errors.Trace(err1) - }) - return dbInfo, tbl, err -} - func updateDDLJob2Table( ctx context.Context, se *sess.Session, diff --git a/pkg/ddl/reorg.go b/pkg/ddl/reorg.go index a190a92ec4183..2f7dad05b6526 100644 --- a/pkg/ddl/reorg.go +++ b/pkg/ddl/reorg.go @@ -802,7 +802,7 @@ func GetTableMaxHandle(ctx *ReorgContext, store kv.Storage, startTS uint64, tbl row := chk.GetRow(0) if tblInfo.IsCommonHandle { pkIdx := tables.FindPrimaryIndex(tblInfo) - maxHandle, err = buildCommonHandleFromChunkRow(time.UTC, tblInfo, pkIdx, handleCols, row) + maxHandle, err = buildCommonHandleFromChunkRow(tbl.UseNewCollate(), time.UTC, tblInfo, pkIdx, handleCols, row) return maxHandle, false, err } return kv.IntHandle(row.GetInt64(0)), false, nil @@ -847,7 +847,7 @@ func buildHandleCols(tbl table.PhysicalTable) []*model.ColumnInfo { return handleCols } -func buildCommonHandleFromChunkRow(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, +func buildCommonHandleFromChunkRow(useNewCollate bool, loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, cols []*model.ColumnInfo, row chunk.Row) (kv.Handle, error) { fieldTypes := make([]*types.FieldType, 0, len(cols)) for _, col := range cols { @@ -857,7 +857,7 @@ func buildCommonHandleFromChunkRow(loc *time.Location, tblInfo *model.TableInfo, tablecodec.TruncateIndexValues(tblInfo, idxInfo, datumRow) var handleBytes []byte - handleBytes, err := codec.EncodeKey(loc, nil, datumRow...) + handleBytes, err := codec.NewEncoder(useNewCollate).EncodeKey(loc, nil, datumRow...) if err != nil { return nil, err } diff --git a/pkg/dxf/importinto/BUILD.bazel b/pkg/dxf/importinto/BUILD.bazel index 1cf63949bd63c..39f544ea50307 100644 --- a/pkg/dxf/importinto/BUILD.bazel +++ b/pkg/dxf/importinto/BUILD.bazel @@ -74,6 +74,7 @@ go_library( "//pkg/util", "//pkg/util/backoff", "//pkg/util/chunk", + "//pkg/util/collate", "//pkg/util/dbterror/exeerrors", "//pkg/util/disttask", "//pkg/util/logutil", diff --git a/pkg/dxf/importinto/conflictedkv/handler.go b/pkg/dxf/importinto/conflictedkv/handler.go index 15bef0c297013..9c06d9ebb95af 100644 --- a/pkg/dxf/importinto/conflictedkv/handler.go +++ b/pkg/dxf/importinto/conflictedkv/handler.go @@ -147,7 +147,7 @@ func (h *BaseHandler) encodeAndHandleRow(ctx context.Context, handle tidbkv.Handle, val []byte) (err error) { tblMeta := h.targetTable.Meta() decodedData, _, err := tables.DecodeRawRowData(h.encoder.SessionCtx.GetExprCtx(), - tblMeta, handle, h.targetTable.Cols(), val) + h.targetTable, handle, h.targetTable.Cols(), val) if err != nil { return errors.Trace(err) } diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index adc4eff13ceec..6619dcf9543ca 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -39,6 +39,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/table/tables" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "go.uber.org/zap" ) @@ -351,7 +352,11 @@ func (*PostProcessSpec) ToSubtaskMeta(planCtx planner.PlanCtx) ([]byte, error) { func buildControllerForPlan(p *LogicalPlan) (*importer.LoadDataController, error) { plan, stmt := &p.Plan, p.Stmt idAlloc := kv.NewPanickingAllocators(plan.TableInfo.SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, plan.TableInfo) + tbl, err := tables.TableFromMetaWithCollate( + plan.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + idAlloc, + plan.TableInfo, + ) if err != nil { return nil, err } diff --git a/pkg/dxf/importinto/task_executor.go b/pkg/dxf/importinto/task_executor.go index 534bc2e830707..b97fe6fcd4e85 100644 --- a/pkg/dxf/importinto/task_executor.go +++ b/pkg/dxf/importinto/task_executor.go @@ -51,6 +51,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool" "github.com/pingcap/tidb/pkg/table/tables" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "github.com/pingcap/tidb/pkg/util/logutil" "go.uber.org/zap" @@ -106,7 +107,11 @@ func getTableImporter( logger *zap.Logger, ) (*importer.TableImporter, error) { idAlloc := kv.NewPanickingAllocators(taskMeta.Plan.TableInfo.SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, taskMeta.Plan.TableInfo) + tbl, err := tables.TableFromMetaWithCollate( + taskMeta.Plan.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + idAlloc, + taskMeta.Plan.TableInfo, + ) if err != nil { return nil, err } @@ -123,9 +128,19 @@ func getTableImporter( } failpoint.Inject("createTableImporterForTest", func() { - failpoint.Return(importer.NewTableImporterForTest(ctx, controller, strconv.FormatInt(taskID, 10), store)) + failpoint.Return(importer.NewTableImporterForTest( + ctx, + controller, + strconv.FormatInt(taskID, 10), + store, + )) }) - return importer.NewTableImporter(ctx, controller, strconv.FormatInt(taskID, 10), store) + return importer.NewTableImporter( + ctx, + controller, + strconv.FormatInt(taskID, 10), + store, + ) } func (s *importStepExecutor) Init(ctx context.Context) (err error) { diff --git a/pkg/executor/admin.go b/pkg/executor/admin.go index d2f73236304df..950c5cf40b052 100644 --- a/pkg/executor/admin.go +++ b/pkg/executor/admin.go @@ -35,7 +35,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/ranger" "github.com/pingcap/tidb/pkg/util/timeutil" @@ -364,7 +363,6 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists idxValLen := len(e.index.Meta().Columns) result.scanRowCount = 0 - useNewCollate := collate.NewCollationEnabled() for { err := srcResult.Next(ctx, e.srcChunk) if err != nil { @@ -406,7 +404,12 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists return nil, err } e.idxValsBufs[result.scanRowCount] = idxVals - rsData := tables.TryGetHandleRestoredDataWrapper(useNewCollate, e.table.Meta(), plannercore.GetCommonHandleDatum(e.handleCols, row), nil, e.index.Meta()) + rsData := tables.TryGetHandleRestoredDataWrapper( + e.table, + plannercore.GetCommonHandleDatum(e.handleCols, row), + nil, + e.index.Meta(), + ) e.recoverRows = append(e.recoverRows, recoverRows{handle: handle, idxVals: idxVals, rsData: rsData, skip: true}) result.scanRowCount++ result.currentHandle = handle diff --git a/pkg/executor/batch_checker.go b/pkg/executor/batch_checker.go index a2d8a8a40b4fa..81209bd6dc1a8 100644 --- a/pkg/executor/batch_checker.go +++ b/pkg/executor/batch_checker.go @@ -300,7 +300,7 @@ func getOldRow(ctx context.Context, sctx sessionctx.Context, txn kv.Transaction, } cols := t.WritableCols() - oldRow, oldRowMap, err := tables.DecodeRawRowData(sctx.GetExprCtx(), t.Meta(), handle, cols, oldValue) + oldRow, oldRowMap, err := tables.DecodeRawRowData(sctx.GetExprCtx(), t, handle, cols, oldValue) if err != nil { return nil, err } diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index 005fe5fc0adc9..ee86cb36c5eb6 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -120,7 +120,7 @@ go_test( embed = [":importer"], flaky = True, race = "on", - shard_count = 41, + shard_count = 42, deps = [ "//br/pkg/mock", "//br/pkg/streamhelper", diff --git a/pkg/executor/importer/chunk_process_testkit_test.go b/pkg/executor/importer/chunk_process_testkit_test.go index c47dc56ef17bb..6a89be5c738f1 100644 --- a/pkg/executor/importer/chunk_process_testkit_test.go +++ b/pkg/executor/importer/chunk_process_testkit_test.go @@ -92,6 +92,7 @@ func TestFileChunkProcess(t *testing.T) { }, &importer.LoadDataController{ ASTArgs: &importer.ASTArgs{}, + Table: table, InsertColumns: table.VisibleCols(), FieldMappings: fieldMappings, }, diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 86fd170494952..2b28a32577cb8 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -335,6 +335,12 @@ type Plan struct { ManualRecovery bool // the keyspace name when submitting this job, only for import-into Keyspace string + // UseNewCollate captures whether the new collation implementation was enabled + // when this import plan's target table snapshot was created. Import execution + // may happen in another keyspace, so key and expression encoding must use this + // captured value instead of the executor process default. Nil means old metadata + // and should fall back to the caller-provided default. + UseNewCollate *bool `json:"use_new_collate,omitempty"` } // GetOnDupKeyMode returns the conflict handling mode. @@ -349,6 +355,21 @@ func (p *Plan) GetOnDupKeyMode() OnDupKeyMode { return p.OnDupKey } +// GetUseNewCollateOrDefault returns the captured new-collation mode, or +// defaultVal for import metadata generated before the field existed. +func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { + if p.UseNewCollate == nil { + return defaultVal + } + return *p.UseNewCollate +} + +// setUseNewCollate stores the new-collation mode captured from the target table +// snapshot. +func (p *Plan) setUseNewCollate(useNewCollate bool) { + p.UseNewCollate = &useNewCollate +} + // ASTArgs is the arguments for ast.LoadDataStmt. // TODO: remove this struct and use the struct which can be serialized. type ASTArgs struct { @@ -553,6 +574,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } + p.setUseNewCollate(tbl.UseNewCollate()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } @@ -1924,6 +1946,7 @@ func createColAssignSimpleExprs( assignments []*ast.Assignment, ctx expression.BuildContext, mu *sync.Mutex, + useNewCollate bool, ) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { if mu != nil { mu.Lock() @@ -1932,7 +1955,7 @@ func createColAssignSimpleExprs( res := make([]expression.Expression, 0, len(assignments)) var allWarnings []contextutil.SQLWarn for _, assign := range assignments { - newExpr, err := expression.BuildSimpleExpr(ctx, assign.Expr) + newExpr, err := expression.BuildSimpleExpr(ctx, assign.Expr, expression.WithUseNewCollate(useNewCollate)) // col assign expr warnings is static, we should generate it for each row processed. // so we save it and clear it here. if ctx.GetEvalCtx().WarningCount() > 0 { @@ -1948,7 +1971,12 @@ func createColAssignSimpleExprs( // CreateColAssignSimpleExprs creates the column assignment expressions using `expression.BuildContext`. func (e *LoadDataController) CreateColAssignSimpleExprs(ctx expression.BuildContext) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { - return createColAssignSimpleExprs(e.ColumnAssignments, ctx, &e.colAssignMu) + return createColAssignSimpleExprs( + e.ColumnAssignments, + ctx, + &e.colAssignMu, + e.Table.UseNewCollate(), + ) } func (e *LoadDataController) getLocalBackendCfg(keyspace, pdAddr, dataDir string) ingestctrl.BackendConfig { diff --git a/pkg/executor/importer/import_test.go b/pkg/executor/importer/import_test.go index b7cfa813fe337..73406210467dc 100644 --- a/pkg/executor/importer/import_test.go +++ b/pkg/executor/importer/import_test.go @@ -105,6 +105,26 @@ func TestInitDefaultOptions(t *testing.T) { require.Equal(t, 5, plan.ThreadCnt) } +func TestPlanUseNewCollate(t *testing.T) { + plan := &Plan{} + require.True(t, plan.GetUseNewCollateOrDefault(true)) + require.False(t, plan.GetUseNewCollateOrDefault(false)) + + plan.setUseNewCollate(false) + require.False(t, plan.GetUseNewCollateOrDefault(true)) + + data, err := json.Marshal(plan) + require.NoError(t, err) + require.Contains(t, string(data), `"use_new_collate":false`) + + var decoded Plan + require.NoError(t, json.Unmarshal(data, &decoded)) + require.False(t, decoded.GetUseNewCollateOrDefault(true)) + + decoded.setUseNewCollate(true) + require.True(t, decoded.GetUseNewCollateOrDefault(false)) +} + // for negative case see TestImportIntoOptionsNegativeCase func TestInitOptionsPositiveCase(t *testing.T) { sctx := mock.NewContext() diff --git a/pkg/executor/importer/sampler.go b/pkg/executor/importer/sampler.go index e5f6bda1cc6ad..073ac642d30a7 100644 --- a/pkg/executor/importer/sampler.go +++ b/pkg/executor/importer/sampler.go @@ -209,7 +209,12 @@ func validateKVSizeSampleConfig(cfg *KVSizeSampleConfig) error { func (s *kvSizeSampler) CreateColAssignSimpleExprs( ctx expression.BuildContext, ) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { - return createColAssignSimpleExprs(s.cfg.ColumnAssignments, ctx, &s.colAssignMu) + return createColAssignSimpleExprs( + s.cfg.ColumnAssignments, + ctx, + &s.colAssignMu, + s.table.UseNewCollate(), + ) } func (s *kvSizeSampler) generateCSVConfig() *config.CSVConfig { @@ -331,7 +336,7 @@ func (s *kvSizeSampler) sampleOneFile( ParquetMeta: file.ParquetMeta, } idAlloc := kv.NewPanickingAllocators(s.table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, s.table.Meta()) + tbl, err := tables.TableFromMetaWithCollate(s.table.UseNewCollate(), idAlloc, s.table.Meta()) if err != nil { return 0, 0, 0, errors.Annotatef(err, "failed to tables.TableFromMeta %s", s.table.Meta().Name) } diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go index a96b4c5337619..655075ac91efd 100644 --- a/pkg/executor/importer/table_import.go +++ b/pkg/executor/importer/table_import.go @@ -91,6 +91,15 @@ var ( defaultMaxEngineSize = int64(5 * config.DefaultBatchSize) ) +func newEncodingTable(e *LoadDataController) (table.Table, error) { + idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) + tbl, err := tables.TableFromMetaWithCollate(e.Table.UseNewCollate(), idAlloc, e.Table.Meta()) + if err != nil { + return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + } + return tbl, nil +} + // Chunk records the chunk information. type Chunk struct { Path string @@ -189,10 +198,9 @@ func NewTableImporter( id string, kvStore tidbkv.Storage, ) (ti *TableImporter, err error) { - idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, e.Table.Meta()) + tbl, err := newEncodingTable(e) if err != nil { - return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + return nil, err } tidbCfg := tidb.GetGlobalConfig() @@ -284,12 +292,16 @@ func (s *storeHelper) GetTiKVCodec() tikv.Codec { var _ ingestctrl.StoreHelper = (*storeHelper)(nil) // NewTableImporterForTest creates a new table importer for test. -func NewTableImporterForTest(ctx context.Context, e *LoadDataController, id string, kvStore tidbkv.Storage) (*TableImporter, error) { +func NewTableImporterForTest( + ctx context.Context, + e *LoadDataController, + id string, + kvStore tidbkv.Storage, +) (*TableImporter, error) { helper := &storeHelper{kvStore: kvStore} - idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, e.Table.Meta()) + tbl, err := newEncodingTable(e) if err != nil { - return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + return nil, err } tidbCfg := tidb.GetGlobalConfig() diff --git a/pkg/expression/expression.go b/pkg/expression/expression.go index a64d9573faec9..aa396e8d5a33a 100644 --- a/pkg/expression/expression.go +++ b/pkg/expression/expression.go @@ -30,6 +30,7 @@ import ( "github.com/pingcap/tidb/pkg/planner/cascades/base" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/generatedexpr" "github.com/pingcap/tidb/pkg/util/size" "github.com/pingcap/tidb/pkg/util/zeropool" @@ -1106,6 +1107,18 @@ func TableInfo2SchemaAndNames(ctx BuildContext, dbName ast.CIStr, tbl *model.Tab // ColumnInfos2ColumnsAndNames converts the ColumnInfo to the *Column and NameSlice. func ColumnInfos2ColumnsAndNames(ctx BuildContext, dbName, tblName ast.CIStr, colInfos []*model.ColumnInfo, tblInfo *model.TableInfo) ([]*Column, types.NameSlice, error) { + return ColumnInfos2ColumnsAndNamesWithCollate(ctx, dbName, tblName, colInfos, tblInfo, collate.NewCollationEnabled()) +} + +// ColumnInfos2ColumnsAndNamesWithCollate converts the ColumnInfo to the *Column +// and NameSlice with a fixed collation mode. +func ColumnInfos2ColumnsAndNamesWithCollate( + ctx BuildContext, + dbName, tblName ast.CIStr, + colInfos []*model.ColumnInfo, + tblInfo *model.TableInfo, + useNewCollate bool, +) ([]*Column, types.NameSlice, error) { columns := make([]*Column, 0, len(colInfos)) names := make([]*types.FieldName, 0, len(colInfos)) for i, col := range colInfos { @@ -1146,7 +1159,10 @@ func ColumnInfos2ColumnsAndNames(ctx BuildContext, dbName, tblName ast.CIStr, co if err != nil { return nil, nil, errors.Trace(err) } - e, err := BuildSimpleExpr(ctx, expr, WithInputSchemaAndNames(mockSchema, names, tblInfo), WithAllowCastArray(true)) + e, err := BuildSimpleExpr(ctx, expr, + WithInputSchemaAndNames(mockSchema, names, tblInfo), + WithAllowCastArray(true), + WithUseNewCollate(useNewCollate)) if err != nil { return nil, nil, errors.Trace(err) } diff --git a/pkg/infoschema/BUILD.bazel b/pkg/infoschema/BUILD.bazel index 7a200c9b43a0a..ff71d94bb0064 100644 --- a/pkg/infoschema/BUILD.bazel +++ b/pkg/infoschema/BUILD.bazel @@ -50,6 +50,7 @@ go_library( "//pkg/types", "//pkg/util", "//pkg/util/chunk", + "//pkg/util/collate", "//pkg/util/dbterror", "//pkg/util/deadlockhistory", "//pkg/util/domainutil", diff --git a/pkg/infoschema/tables.go b/pkg/infoschema/tables.go index 3072d29e07c87..1c80ee1a4b6a3 100644 --- a/pkg/infoschema/tables.go +++ b/pkg/infoschema/tables.go @@ -55,6 +55,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/deadlockhistory" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/logutil" @@ -2647,6 +2648,12 @@ func (it *infoschemaTable) Meta() *model.TableInfo { return it.meta } +// UseNewCollate implements table.Table UseNewCollate interface. Info schema +// tables are not persisted user tables, so they use the current process default. +func (it *infoschemaTable) UseNewCollate() bool { + return collate.NewCollationEnabled() +} + // GetPhysicalID implements table.Table GetPhysicalID interface. func (it *infoschemaTable) GetPhysicalID() int64 { return it.meta.ID @@ -2745,6 +2752,12 @@ func (vt *VirtualTable) Meta() *model.TableInfo { return nil } +// UseNewCollate implements table.Table UseNewCollate interface. Virtual tables +// are not persisted user tables, so they use the current process default. +func (vt *VirtualTable) UseNewCollate() bool { + return collate.NewCollationEnabled() +} + // GetPhysicalID implements table.Table GetPhysicalID interface. func (vt *VirtualTable) GetPhysicalID() int64 { return 0 diff --git a/pkg/lightning/backend/kv/base.go b/pkg/lightning/backend/kv/base.go index 05ac73a7eaf96..b05c4dc784731 100644 --- a/pkg/lightning/backend/kv/base.go +++ b/pkg/lightning/backend/kv/base.go @@ -167,7 +167,7 @@ func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) { } // collect expressions for evaluating stored generated columns - genCols, err := CollectGeneratedColumns(se, meta, cols) + genCols, err := CollectGeneratedColumns(se, config.Table) if err != nil { return nil, errors.Annotate(err, "failed to parse generated column expressions") } diff --git a/pkg/lightning/backend/kv/kv2sql.go b/pkg/lightning/backend/kv/kv2sql.go index 0a535166b1902..fa589c8fb2829 100644 --- a/pkg/lightning/backend/kv/kv2sql.go +++ b/pkg/lightning/backend/kv/kv2sql.go @@ -54,7 +54,7 @@ func (t *TableKVDecoder) DecodeHandleFromIndex(indexInfo *model.IndexInfo, key, // DecodeRawRowData decodes raw row data into a datum slice and a (columnID:columnValue) map. func (t *TableKVDecoder) DecodeRawRowData(h kv.Handle, value []byte) ([]types.Datum, map[int64]types.Datum, error) { - return tables.DecodeRawRowData(t.se.GetExprCtx(), t.tbl.Meta(), h, t.tbl.Cols(), value) + return tables.DecodeRawRowData(t.se.GetExprCtx(), t.tbl, h, t.tbl.Cols(), value) } // DecodeRawRowDataAsStr decodes raw row data into a string. @@ -134,8 +134,7 @@ func NewTableKVDecoder( return nil, err } - cols := tbl.Cols() - genCols, err := CollectGeneratedColumns(se, tbl.Meta(), cols) + genCols, err := CollectGeneratedColumns(se, tbl) if err != nil { return nil, err } diff --git a/pkg/lightning/backend/kv/sql2kv.go b/pkg/lightning/backend/kv/sql2kv.go index 31666a97fa947..140c8eadccd30 100644 --- a/pkg/lightning/backend/kv/sql2kv.go +++ b/pkg/lightning/backend/kv/sql2kv.go @@ -68,7 +68,9 @@ func NewTableKVEncoder( // CollectGeneratedColumns collects all expressions required to evaluate the // results of all generated columns. The returning slice is in evaluation order. -func CollectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.Column) ([]GeneratedCol, error) { +func CollectGeneratedColumns(se *Session, tbl table.Table) ([]GeneratedCol, error) { + meta := tbl.Meta() + cols := tbl.Cols() hasGenCol := false for _, col := range cols { if col.GeneratedExpr != nil { @@ -112,6 +114,7 @@ func CollectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.C col.GeneratedExpr.Internal(), expression.WithInputSchemaAndNames(schema, names, meta), expression.WithAllowCastArray(true), + expression.WithUseNewCollate(tbl.UseNewCollate()), ) if err != nil { return nil, err diff --git a/pkg/meta/model/job_test.go b/pkg/meta/model/job_test.go index 5739b332e3007..95d8f666e506b 100644 --- a/pkg/meta/model/job_test.go +++ b/pkg/meta/model/job_test.go @@ -117,6 +117,26 @@ func TestJobCodec(t *testing.T) { require.Equal(t, int64(3), job.GetRowCount()) } +func TestDDLReorgMetaUseNewCollate(t *testing.T) { + meta := &DDLReorgMeta{} + require.True(t, meta.GetUseNewCollateOrDefault(true)) + require.False(t, meta.GetUseNewCollateOrDefault(false)) + + meta.setUseNewCollate(false) + require.False(t, meta.GetUseNewCollateOrDefault(true)) + + data, err := json.Marshal(meta) + require.NoError(t, err) + require.Contains(t, string(data), `"use_new_collate":false`) + + var decoded DDLReorgMeta + require.NoError(t, json.Unmarshal(data, &decoded)) + require.False(t, decoded.GetUseNewCollateOrDefault(true)) + + decoded.setUseNewCollate(true) + require.True(t, decoded.GetUseNewCollateOrDefault(false)) +} + func TestLocation(t *testing.T) { // test offset = 0 loc := &TimeZoneLocation{} diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index e904c903dfc33..76d7e34cf729b 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -95,6 +95,12 @@ type DDLReorgMeta struct { MaxNodeCount int `json:"max_node_count"` AnalyzeState int8 `json:"analyze_state"` Stage ReorgStage `json:"stage"` + // UseNewCollate captures whether the new collation implementation was enabled + // when this reorg task's persisted table snapshot was created. Reorg execution + // may happen in another keyspace, so key and expression encoding must use this + // captured value instead of the executor process default. Nil means old metadata + // and should fall back to the caller-provided default. + UseNewCollate *bool `json:"use_new_collate,omitempty"` // These two variables are used to control the concurrency and batch size of the reorganization process. // They can be adjusted dynamically through `admin alter ddl jobs` command. // Note: Don't get or set these two variables directly, use the functions instead. @@ -151,6 +157,21 @@ func (dm *DDLReorgMeta) SetMaxWriteSpeed(maxWriteSpeed int) { dm.MaxWriteSpeed.Store(int64(maxWriteSpeed)) } +// GetUseNewCollateOrDefault returns the captured new-collation mode, or +// defaultVal for reorg metadata generated before the field existed. +func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { + if dm.UseNewCollate == nil { + return defaultVal + } + return *dm.UseNewCollate +} + +// setUseNewCollate stores the new-collation mode captured from the persisted +// table snapshot. +func (dm *DDLReorgMeta) setUseNewCollate(useNewCollate bool) { + dm.UseNewCollate = &useNewCollate +} + const ( // ReorgMetaVersion0 is the minimum version of DDLReorgMeta. ReorgMetaVersion0 = int64(0) diff --git a/pkg/table/table.go b/pkg/table/table.go index 575f2e18c7a24..7fd8fc9d11ca7 100644 --- a/pkg/table/table.go +++ b/pkg/table/table.go @@ -459,6 +459,14 @@ type Table interface { // Meta returns TableInfo. Meta() *model.TableInfo + // UseNewCollate returns whether the new collation implementation should be used + // when encoding persisted keys or expressions for this table. + // + // For persisted user tables, this value must come from the table snapshot used + // to build the task. Virtual or non-persisted table implementations may use the + // current process default. + UseNewCollate() bool + // Type returns the type of table Type() Type diff --git a/pkg/table/tables/BUILD.bazel b/pkg/table/tables/BUILD.bazel index eb6dac9730867..9775792b76383 100644 --- a/pkg/table/tables/BUILD.bazel +++ b/pkg/table/tables/BUILD.bazel @@ -80,7 +80,7 @@ go_test( ], embed = [":tables"], flaky = True, - shard_count = 43, + shard_count = 44, deps = [ "//pkg/ddl", "//pkg/domain", diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index e3b0e3076301f..7e83ba2fa251e 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -392,8 +392,7 @@ func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.C if err != nil { return nil, nil, err } - useNewCollate := table.encoder.UseNewCollate() - rsData := TryGetHandleRestoredDataWrapper(useNewCollate, table.meta, rowToInsert, nil, indexInfo) + rsData := TryGetHandleRestoredDataWrapper(table, rowToInsert, nil, indexInfo) value, err := index.GenIndexValue(errctx.StrictNoWarningContext, tc.Location(), distinct, false, indexedValues, handle, rsData, nil) if err != nil { return nil, nil, err diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index 8ee3895538d3a..0f2cd754e9a56 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -126,7 +126,7 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part return initEvalBuffer(ret) }, } - if err := initTableIndices(&ret.TableCommon); err != nil { + if err := ret.initTableIndices(); err != nil { return nil, errors.Trace(err) } origIndices := ret.meta.Indices @@ -181,8 +181,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part } else { tblInfo.Indices = origIndices } - err := initTableCommonWithIndices(&t.TableCommon, tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints) - if err != nil { + t.TableCommon = newTableCommon(tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints, tbl.encoder.UseNewCollate()) + if err := t.initTableIndices(); err != nil { return nil, errors.Trace(err) } t.table = ret @@ -273,8 +273,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part func initPartition(t *partitionedTable, def model.PartitionDefinition) (*partition, error) { var newPart partition - err := initTableCommonWithIndices(&newPart.TableCommon, t.meta, def.ID, t.Columns, t.allocs, t.Constraints) - if err != nil { + newPart.TableCommon = newTableCommon(t.meta, def.ID, t.Columns, t.allocs, t.Constraints, t.encoder.UseNewCollate()) + if err := newPart.initTableIndices(); err != nil { return nil, err } newPart.table = t @@ -1675,8 +1675,7 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro if err != nil { return nil, err } - var tc TableCommon - tc.initTableCommon(tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc := newTableCommon(tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints, t.UseNewCollate()) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index da525127a579a..69ea903197bb3 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -141,8 +141,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { if err != nil { return nil } - var t TableCommon - t.initTableCommon(tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) + t := newTableCommon(tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints, collate.NewCollationEnabled()) if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { ret, err := newCachedTable(&t) if err != nil { @@ -151,7 +150,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { return ret } if tblInfo.GetPartitionInfo() == nil { - if err := initTableIndices(&t); err != nil { + if err := t.initTableIndices(); err != nil { return nil } return &t @@ -166,6 +165,12 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { // TableFromMeta creates a Table instance from model.TableInfo. func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Table, error) { + return TableFromMetaWithCollate(collate.NewCollationEnabled(), allocs, tblInfo) +} + +// TableFromMetaWithCollate creates a Table instance from model.TableInfo with a +// fixed new-collation mode for persisted key and expression encoding. +func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Table, error) { if tblInfo.State == model.StateNone { return nil, table.ErrTableStateCantNone.GenWithStackByArgs(tblInfo.Name) } @@ -215,10 +220,9 @@ func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Ta if err != nil { return nil, err } - var t TableCommon - t.initTableCommon(tblInfo, tblInfo.ID, columns, allocs, constraints) + t := newTableCommon(tblInfo, tblInfo.ID, columns, allocs, constraints, useNewCollate) if tblInfo.GetPartitionInfo() == nil { - if err := initTableIndices(&t); err != nil { + if err := t.initTableIndices(); err != nil { return nil, err } if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { @@ -241,25 +245,27 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, return expr, nil } -// initTableCommon initializes a TableCommon struct. -func (t *TableCommon) initTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { - t.tableID = tblInfo.ID - t.physicalTableID = physicalTableID - t.allocs = allocs - t.meta = tblInfo - t.Columns = cols - t.Constraints = constraints - t.recordPrefix = tablecodec.GenTableRecordPrefix(physicalTableID) - t.indexPrefix = tablecodec.GenTableIndexPrefix(physicalTableID) - t.encoder = codec.NewEncoder(collate.NewCollationEnabled()) +func newTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint, useNewCollate bool) TableCommon { + t := TableCommon{ + tableID: tblInfo.ID, + physicalTableID: physicalTableID, + allocs: allocs, + meta: tblInfo, + Columns: cols, + Constraints: constraints, + recordPrefix: tablecodec.GenTableRecordPrefix(physicalTableID), + indexPrefix: tablecodec.GenTableIndexPrefix(physicalTableID), + encoder: codec.NewEncoder(useNewCollate), + } if tblInfo.IsSequence() { t.sequence = &sequenceCommon{meta: tblInfo.Sequence} } t.ResetColumnsCache() + return t } // initTableIndices initializes the indices of the TableCommon. -func initTableIndices(t *TableCommon) error { +func (t *TableCommon) initTableIndices() error { tblInfo := t.meta for _, idxInfo := range tblInfo.Indices { if idxInfo.State == model.StateNone { @@ -310,11 +316,6 @@ func asIndex(idx table.Index) *index { return idx.(*index) } -func initTableCommonWithIndices(t *TableCommon, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - t.initTableCommon(tblInfo, physicalTableID, cols, allocs, constraints) - return initTableIndices(t) -} - // Indices implements table.Table Indices interface. func (t *TableCommon) Indices() []table.Index { return t.indices @@ -1034,7 +1035,7 @@ func (t *TableCommon) addIndices(sctx table.MutateContext, recordID kv.Handle, r } dupErr = kv.GenKeyExistsErr(colStrVals, fmt.Sprintf("%s.%s", v.TableMeta().Name.String(), v.Meta().Name.String())) } - rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, r, nil, v.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t, r, nil, v.Meta()) if dupHandle, err := asIndex(v).create(sctx, txn, indexVals, recordID, rsData, false, opt); err != nil { if kv.ErrKeyExists.Equal(err) { return dupHandle, dupErr @@ -1059,7 +1060,7 @@ func RowWithCols(t table.Table, ctx sessionctx.Context, h kv.Handle, cols []*tab if err != nil { return nil, err } - v, _, err := DecodeRawRowData(ctx.GetExprCtx(), t.Meta(), h, cols, value) + v, _, err := DecodeRawRowData(ctx.GetExprCtx(), t, h, cols, value) if err != nil { return nil, err } @@ -1079,8 +1080,10 @@ func containFullColInHandle(meta *model.TableInfo, col *table.Column) (containFu } // DecodeRawRowData decodes raw row data into a datum slice and a (columnID:columnValue) map. -func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.Handle, cols []*table.Column, +func DecodeRawRowData(ctx expression.BuildContext, tbl table.Table, h kv.Handle, cols []*table.Column, value []byte) ([]types.Datum, map[int64]types.Datum, error) { + meta := tbl.Meta() + useNewCollate := tbl.UseNewCollate() v := make([]types.Datum, len(cols)) colTps := make(map[int64]*types.FieldType, len(cols)) prefixCols := make(map[int64]struct{}) @@ -1096,7 +1099,7 @@ func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.H } continue } - if col.IsCommonHandleColumn(meta) && !types.NeedRestoredData(&col.FieldType) { + if col.IsCommonHandleColumn(meta) && !types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate) { if containFullCol, idxInHandle := containFullColInHandle(meta, col); containFullCol { dtBytes := h.EncodedCol(idxInHandle) _, dt, err := codec.DecodeOne(dtBytes) @@ -1123,7 +1126,8 @@ func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.H if col == nil { continue } - if col.IsPKHandleColumn(meta) || (col.IsCommonHandleColumn(meta) && !types.NeedRestoredData(&col.FieldType)) { + if col.IsPKHandleColumn(meta) || + (col.IsCommonHandleColumn(meta) && !types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate)) { if _, isPrefix := prefixCols[col.ID]; !isPrefix { continue } @@ -1314,7 +1318,7 @@ func (t *TableCommon) removeRowIndices(ctx table.MutateContext, txn kv.Transacti } func (t *TableCommon) buildIndexForRow(ctx table.MutateContext, h kv.Handle, vals []types.Datum, newData []types.Datum, idx *index, txn kv.Transaction, untouched bool, opt *table.CreateIdxOpt) error { - rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, newData, nil, idx.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t, newData, nil, idx.Meta()) if _, err := idx.create(ctx, txn, vals, h, rsData, untouched, opt); err != nil { if kv.ErrKeyExists.Equal(err) { // Make error message consistent with MySQL. @@ -1520,6 +1524,11 @@ func (t *TableCommon) Type() table.Type { return table.NormalTable } +// UseNewCollate implements table.Table UseNewCollate interface. +func (t *TableCommon) UseNewCollate() bool { + return t.encoder.UseNewCollate() +} + func (t *TableCommon) canSkip(col *table.Column, value *types.Datum) bool { return CanSkip(t.encoder.UseNewCollate(), t.Meta(), col, value) } @@ -1777,8 +1786,10 @@ func (t *TableCommon) GetSequenceCommon() *sequenceCommon { // TryGetHandleRestoredDataWrapper tries to get the restored data for handle if // needed. The argument can be a slice or a map. -func TryGetHandleRestoredDataWrapper(useNewCollate bool, tblInfo *model.TableInfo, +func TryGetHandleRestoredDataWrapper(tbl table.Table, row []types.Datum, rowMap map[int64]types.Datum, idx *model.IndexInfo) []types.Datum { + useNewCollate := tbl.UseNewCollate() + tblInfo := tbl.Meta() if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { return nil } diff --git a/pkg/table/tables/tables_test.go b/pkg/table/tables/tables_test.go index 3e4fd0eca2041..41543f3a913bc 100644 --- a/pkg/table/tables/tables_test.go +++ b/pkg/table/tables/tables_test.go @@ -408,6 +408,29 @@ func TestTableFromMeta(t *testing.T) { require.Error(t, err) } +func TestTableFromMetaWithCollateUsesFixedMode(t *testing.T) { + tblInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + State: model.StatePublic, + Columns: []*model.ColumnInfo{ + { + ID: 1, + Name: ast.NewCIStr("a"), + Offset: 0, + State: model.StatePublic, + FieldType: *types.NewFieldType(mysql.TypeVarchar), + }, + }, + } + + for _, useNewCollate := range []bool{false, true} { + tbl, err := tables.TableFromMetaWithCollate(useNewCollate, autoid.NewAllocators(false), tblInfo) + require.NoError(t, err) + require.Equal(t, useNewCollate, tbl.UseNewCollate()) + } +} + func TestHiddenColumn(t *testing.T) { store, dom := testkit.CreateMockStoreAndDomain(t) tk := testkit.NewTestKit(t, store) diff --git a/tests/realtikvtest/addindextest1/BUILD.bazel b/tests/realtikvtest/addindextest1/BUILD.bazel index f0673be31315e..6a82811448c7d 100644 --- a/tests/realtikvtest/addindextest1/BUILD.bazel +++ b/tests/realtikvtest/addindextest1/BUILD.bazel @@ -9,7 +9,7 @@ go_test( "main_test.go", ], flaky = True, - shard_count = 14, + shard_count = 15, deps = [ "//pkg/config", "//pkg/config/kerneltype", @@ -30,6 +30,7 @@ go_test( "//pkg/store", "//pkg/testkit", "//pkg/testkit/testfailpoint", + "//pkg/util/collate", "//pkg/util/dbterror", "//tests/realtikvtest", "@com_github_pingcap_errors//:errors", diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index f03e897f20e97..c16574fcb7ac9 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -17,15 +17,19 @@ package addindextest import ( "fmt" "strconv" + "sync/atomic" "testing" + "github.com/pingcap/tidb/pkg/config" "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/ddl" "github.com/pingcap/tidb/pkg/ddl/schemaver" "github.com/pingcap/tidb/pkg/keyspace" + "github.com/pingcap/tidb/pkg/meta/model" kvstore "github.com/pingcap/tidb/pkg/store" "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/testkit/testfailpoint" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/tests/realtikvtest" "github.com/stretchr/testify/require" ) @@ -60,6 +64,183 @@ func TestAddIndexOnSystemTable(t *testing.T) { tk.MustQuery(taskQuerySQL).Check(testkit.Rows("0")) } +func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("This test is only for nextgen kernel, skip it in classic kernel") + } + restore := config.RestoreFunc() + t.Cleanup(restore) + config.UpdateGlobal(func(conf *config.Config) { + conf.Experimental.AllowsExpressionIndex = true + }) + + originNewCollationEnabled := collate.NewCollationEnabled() + t.Cleanup(func() { + collate.SetNewCollationEnabledForTest(originNewCollationEnabled) + }) + + const userKeyspace = "keyspacecollate" + runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ + keyspace.System: true, + userKeyspace: false, + }, userKeyspace) + userStore := runtimes[userKeyspace].Store + tk := testkit.NewTestKit(t, userStore) + tk.MustQuery(`select variable_value from mysql.tidb where variable_name = 'new_collation_enabled'`). + Check(testkit.Rows("False")) + require.False(t, collate.NewCollationEnabled()) + + var backfillInitCnt atomic.Int64 + testfailpoint.Enable( + t, + "github.com/pingcap/tidb/pkg/ddl/overrideDefaultUseNewCollateForBackfillStep", + "return(true)", + ) + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/ddl/afterResolveUserTableNewCollateForBackfillStep", + func(job *model.Job, defaultUseNewCollate bool, useNewCollate bool) { + require.False(t, job.ReorgMeta.GetUseNewCollateOrDefault(true)) + require.True(t, defaultUseNewCollate) + require.False(t, useNewCollate) + require.False(t, collate.NewCollationEnabled()) + backfillInitCnt.Add(1) + }, + ) + + tk.PrepareDB("crossks_collate") + + cases := []struct { + name string + table string + setupSQL []string + addIndexSQL []string + indexes []string + dmlSQL []string + }{ + { + name: "clustered varchar primary key and secondary varchar index", + table: "t_varchar_pk", + setupSQL: []string{ + "drop table if exists t_varchar_pk", + `create table t_varchar_pk ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered + )`, + "insert into t_varchar_pk values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_varchar_pk add index idx_fk(fk)", + }, + indexes: []string{"idx_fk"}, + dmlSQL: []string{ + "insert into t_varchar_pk values ('ddd', 'dbc')", + "update t_varchar_pk set fk = 'updated' where id = 'ddd'", + "delete from t_varchar_pk where id = 'ddd'", + }, + }, + { + name: "composite clustered primary key with varchar part and secondary int index", + table: "t_composite_varchar_pk", + setupSQL: []string{ + "drop table if exists t_composite_varchar_pk", + `create table t_composite_varchar_pk ( + id1 varchar(32) collate utf8mb4_general_ci, + id2 int, + fk int, + primary key (id1, id2) clustered + )`, + "insert into t_composite_varchar_pk values ('ax', 1, 10), ('by', 2, 20), ('cz', 3, 30)", + }, + addIndexSQL: []string{ + "alter table t_composite_varchar_pk add index idx_fk(fk)", + }, + indexes: []string{"idx_fk"}, + dmlSQL: []string{ + "insert into t_composite_varchar_pk values ('dw', 4, 40)", + "update t_composite_varchar_pk set fk = 41 where id1 = 'dw' and id2 = 4", + "delete from t_composite_varchar_pk where id1 = 'dw' and id2 = 4", + }, + }, + { + name: "generated columns with string transformations", + table: "t_add_generated_column_index", + setupSQL: []string{ + "drop table if exists t_add_generated_column_index", + `create table t_add_generated_column_index ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + g_lower varchar(32) generated always as (lower(raw)) virtual, + g_upper varchar(32) generated always as (upper(raw)) virtual, + g_concat varchar(80) generated always as (concat(id, ':', raw)) virtual, + g_substr varchar(32) generated always as (substr(raw, 1, 2)) virtual, + primary key (id) clustered + )`, + "insert into t_add_generated_column_index(id, raw) values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_add_generated_column_index add index idx_g_lower(g_lower)", + "alter table t_add_generated_column_index add index idx_g_upper(g_upper)", + "alter table t_add_generated_column_index add index idx_g_concat(g_concat)", + "alter table t_add_generated_column_index add index idx_g_substr(g_substr)", + }, + indexes: []string{"idx_g_lower", "idx_g_upper", "idx_g_concat", "idx_g_substr"}, + dmlSQL: []string{ + "insert into t_add_generated_column_index(id, raw) values ('ddd', 'dbc')", + "update t_add_generated_column_index set raw = 'updated' where id = 'ddd'", + "delete from t_add_generated_column_index where id = 'ddd'", + }, + }, + { + name: "expression indexes with string transformations", + table: "t_add_expression_index", + setupSQL: []string{ + "drop table if exists t_add_expression_index", + `create table t_add_expression_index ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered + )`, + "insert into t_add_expression_index values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_add_expression_index add index idx_lower ((lower(raw)))", + "alter table t_add_expression_index add index idx_upper ((upper(raw)))", + "alter table t_add_expression_index add index idx_concat ((concat(id, ':', raw)))", + "alter table t_add_expression_index add index idx_substr ((substr(raw, 1, 2)))", + }, + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + dmlSQL: []string{ + "insert into t_add_expression_index values ('ddd', 'dbc')", + "update t_add_expression_index set raw = 'updated' where id = 'ddd'", + "delete from t_add_expression_index where id = 'ddd'", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + collate.SetNewCollationEnabledForTest(false) + for _, sql := range tc.setupSQL { + tk.MustExec(sql) + } + for _, sql := range tc.addIndexSQL { + before := backfillInitCnt.Load() + collate.SetNewCollationEnabledForTest(false) + tk.MustExec(sql) + require.Greater(t, backfillInitCnt.Load(), before) + collate.SetNewCollationEnabledForTest(false) + } + checkTableAndIndexes(tk, tc.table, tc.indexes, "3") + for _, sql := range tc.dmlSQL { + tk.MustExec(sql) + } + checkTableAndIndexes(tk, tc.table, tc.indexes, "3") + }) + } +} + func TestCrossKSInfoSchemaSync(t *testing.T) { if kerneltype.IsClassic() { t.Skip("This test is only for nextgen kernel, skip it in classic kernel") @@ -151,3 +332,12 @@ func TestCrossKSInfoSchemaSync(t *testing.T) { require.EqualValues(t, 2, sum.AssumedServerCount) }) } + +func checkTableAndIndexes(tk *testkit.TestKit, tableName string, indexes []string, expectedCount string) { + tk.MustExec("admin check table " + tableName) + tk.MustQuery("select count(*) from " + tableName).Check(testkit.Rows(expectedCount)) + for _, indexName := range indexes { + tk.MustQuery(fmt.Sprintf("select count(*) from %s force index(%s)", tableName, indexName)). + Check(testkit.Rows(expectedCount)) + } +} diff --git a/tests/realtikvtest/configs/next-gen/pd.toml b/tests/realtikvtest/configs/next-gen/pd.toml index 578cfbf2aaab4..07569ded421b1 100644 --- a/tests/realtikvtest/configs/next-gen/pd.toml +++ b/tests/realtikvtest/configs/next-gen/pd.toml @@ -1,3 +1,3 @@ # list all the keyspace names here, each keyspace represents a tenant. [keyspace] -pre-alloc = ["keyspace1", "keyspace2"] +pre-alloc = ["keyspace1", "keyspace2", "keyspacecollate"] diff --git a/tests/realtikvtest/importintotest3/BUILD.bazel b/tests/realtikvtest/importintotest3/BUILD.bazel index 1eddd669eecda..c2262524cc068 100644 --- a/tests/realtikvtest/importintotest3/BUILD.bazel +++ b/tests/realtikvtest/importintotest3/BUILD.bazel @@ -13,7 +13,7 @@ go_test( ], flaky = True, race = "on", - shard_count = 2, + shard_count = 3, deps = [ "//pkg/config", "//pkg/config/deploymode", @@ -24,6 +24,7 @@ go_test( "//pkg/dxf/framework/testutil", "//pkg/dxf/importinto", "//pkg/executor/importer", + "//pkg/keyspace", "//pkg/kv", "//pkg/lightning/mydump", "//pkg/objstore", @@ -33,6 +34,7 @@ go_test( "//pkg/store/driver/error", "//pkg/testkit", "//pkg/testkit/testfailpoint", + "//pkg/util/collate", "//tests/realtikvtest", "@com_github_fsouza_fake_gcs_server//fakestorage", "@com_github_golang_snappy//:snappy", diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 0f32ef77ad522..75a287d3a5ca7 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -20,18 +20,23 @@ import ( "fmt" "strconv" "strings" + "sync/atomic" "testing" "time" "github.com/pingcap/tidb/pkg/config/kerneltype" + "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/dxf/framework/taskexecutor/execute" "github.com/pingcap/tidb/pkg/dxf/importinto" "github.com/pingcap/tidb/pkg/executor/importer" + "github.com/pingcap/tidb/pkg/keyspace" "github.com/pingcap/tidb/pkg/objstore" plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/sessionctx/vardef" kvstore "github.com/pingcap/tidb/pkg/store" "github.com/pingcap/tidb/pkg/testkit" + "github.com/pingcap/tidb/pkg/testkit/testfailpoint" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/tests/realtikvtest" "github.com/stretchr/testify/require" ) @@ -128,3 +133,490 @@ func TestOnUserKeyspace(t *testing.T) { require.NoError(t, json.Unmarshal([]byte(rs[0][0].(string)), summary)) require.EqualValues(t, rowCount, summary.ImportedRows) } + +func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("only runs in nextgen kernel") + } + + originNewCollationEnabled := collate.NewCollationEnabled() + t.Cleanup(func() { + collate.SetNewCollationEnabledForTest(originNewCollationEnabled) + }) + + const userKeyspace = "keyspacecollate" + runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ + keyspace.System: true, + userKeyspace: false, + }, userKeyspace) + userStore := runtimes[userKeyspace].Store + userTK := testkit.NewTestKit(t, userStore) + userTK.MustQuery(`select variable_value from mysql.tidb where variable_name = 'new_collation_enabled'`). + Check(testkit.Rows("False")) + require.False(t, collate.NewCollationEnabled()) + + ctx := context.Background() + s3Args := "access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2f0.0.0.0%3a9000" + objStore, err := objstore.NewFromURL(ctx, fmt.Sprintf("s3://next-gen-test/collate-data?%s", s3Args)) + require.NoError(t, err) + t.Cleanup(func() { + objStore.Close() + }) + + var taskSubmitCnt atomic.Int64 + var taskRefreshCnt atomic.Int64 + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", + func(*int, *proto.ExtraParams) { + collate.SetNewCollationEnabledForTest(true) + taskSubmitCnt.Add(1) + }, + ) + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/afterRefreshTask", + func(task *proto.Task) { + if task == nil || task.Type != proto.ImportInto || !strings.HasPrefix(task.Key, userKeyspace+"/ImportInto/") { + return + } + var taskMeta importinto.TaskMeta + require.NoError(t, json.Unmarshal(task.Meta, &taskMeta)) + require.NotNil(t, taskMeta.Plan.UseNewCollate) + require.False(t, *taskMeta.Plan.UseNewCollate) + taskRefreshCnt.Add(1) + }, + ) + + prepareAndUseDB("cross_ks_collate", userTK) + + cases := []struct { + name string + table string + setupSQL []string + fileName string + fileData string + importSQL string + indexes []string + postDMLSQL []string + }{ + { + name: "clustered varchar primary key and secondary varchar index", + table: "trigger_varchar_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_varchar_idx", + `create table trigger_varchar_pk_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_varchar_pk_varchar_idx.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_varchar_pk_varchar_idx(@1,id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_varchar_idx values ('ddd', 'dbc')", + "update trigger_varchar_pk_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from trigger_varchar_pk_varchar_idx where id = 'ddd'", + }, + }, + { + name: "clustered varchar primary key and secondary int index", + table: "trigger_varchar_pk_int_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_int_idx", + `create table trigger_varchar_pk_int_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk int, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_varchar_pk_int_idx.csv", + fileData: "10,aaa,x\n20,bbb,x\n30,ccc,x\n", + importSQL: "import into trigger_varchar_pk_int_idx(fk,id,@3) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_int_idx values ('ddd', 40)", + "update trigger_varchar_pk_int_idx set fk = 41 where id = 'ddd'", + "delete from trigger_varchar_pk_int_idx where id = 'ddd'", + }, + }, + { + name: "composite clustered primary key with varchar part and secondary int index", + table: "trigger_composite_varchar_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_varchar_int_pk_int_idx", + `create table trigger_composite_varchar_int_pk_int_idx ( + id1 varchar(32) collate utf8mb4_general_ci, + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_varchar_int_pk_int_idx.csv", + fileData: "1,10,aaa\n2,20,bbb\n3,30,ccc\n", + importSQL: "import into trigger_composite_varchar_int_pk_int_idx(id2,fk,id1) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_varchar_int_pk_int_idx values ('ddd', 4, 40)", + "update trigger_composite_varchar_int_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 4", + "delete from trigger_composite_varchar_int_pk_int_idx where id1 = 'ddd' and id2 = 4", + }, + }, + { + name: "composite clustered int primary key and secondary varchar index", + table: "trigger_composite_int_int_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_int_int_pk_varchar_idx", + `create table trigger_composite_int_int_pk_varchar_idx ( + id1 int, + id2 int, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_int_int_pk_varchar_idx.csv", + fileData: "1,10,aaa\n2,20,bbb\n3,30,ccc\n", + importSQL: "import into trigger_composite_int_int_pk_varchar_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_int_int_pk_varchar_idx values (4, 40, 'ddd')", + "update trigger_composite_int_int_pk_varchar_idx set fk = 'updated' where id1 = 4 and id2 = 40", + "delete from trigger_composite_int_int_pk_varchar_idx where id1 = 4 and id2 = 40", + }, + }, + { + name: "composite char primary key and secondary varchar index", + table: "trigger_composite_char_char_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_char_char_pk_varchar_idx", + `create table trigger_composite_char_char_pk_varchar_idx ( + id1 char(32) collate utf8mb4_general_ci, + id2 char(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_char_char_pk_varchar_idx.csv", + fileData: "aaa,ax,ay\nbbb,bx,by\nccc,cx,cy\n", + importSQL: "import into trigger_composite_char_char_pk_varchar_idx(fk,id1,id2) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_char_char_pk_varchar_idx values ('dx', 'dy', 'ddd')", + "update trigger_composite_char_char_pk_varchar_idx set fk = 'updated' where id1 = 'dx' and id2 = 'dy'", + "delete from trigger_composite_char_char_pk_varchar_idx where id1 = 'dx' and id2 = 'dy'", + }, + }, + { + name: "clustered varchar primary key and prefix secondary varchar index", + table: "trigger_varchar_pk_prefix_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_prefix_varchar_idx", + `create table trigger_varchar_pk_prefix_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk_prefix (fk(2)) + )`, + }, + fileName: "trigger_varchar_pk_prefix_varchar_idx.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_varchar_pk_prefix_varchar_idx(@1,id,fk) from '%s'", + indexes: []string{"idx_fk_prefix"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_prefix_varchar_idx values ('ddd', 'dbc')", + "update trigger_varchar_pk_prefix_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from trigger_varchar_pk_prefix_varchar_idx where id = 'ddd'", + }, + }, + { + name: "clustered varchar primary key and secondary varchar index with extra payload", + table: "trigger_record_ok_index_bad_extra_payload", + setupSQL: []string{ + "drop table if exists trigger_record_ok_index_bad_extra_payload", + `create table trigger_record_ok_index_bad_extra_payload ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + payload varchar(32) collate utf8mb4_general_ci default 'payload', + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_record_ok_index_bad_extra_payload.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_record_ok_index_bad_extra_payload(@1,id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_record_ok_index_bad_extra_payload(id, fk) values ('ddd', 'dbc')", + "update trigger_record_ok_index_bad_extra_payload set fk = 'updated' where id = 'ddd'", + "delete from trigger_record_ok_index_bad_extra_payload where id = 'ddd'", + }, + }, + { + name: "generated columns with string transformations", + table: "t_import_generated", + setupSQL: []string{ + "drop table if exists t_import_generated", + `create table t_import_generated ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + g_lower varchar(32) generated always as (lower(raw)) stored, + g_upper varchar(32) generated always as (upper(raw)) stored, + g_concat varchar(80) generated always as (concat(id, ':', raw)) stored, + g_substr varchar(32) generated always as (substr(raw, 1, 2)) stored, + primary key (id) clustered, + key idx_lower (g_lower), + key idx_upper (g_upper), + key idx_concat (g_concat), + key idx_substr (g_substr) + )`, + }, + fileName: "t_import_generated.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into t_import_generated(@1,id,raw) from '%s'", + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + postDMLSQL: []string{ + "insert into t_import_generated(id, raw) values ('ddd', 'dbc')", + "update t_import_generated set raw = 'updated' where id = 'ddd'", + "delete from t_import_generated where id = 'ddd'", + }, + }, + { + name: "assignment expressions with string transformations", + table: "t_import_assignment", + setupSQL: []string{ + "drop table if exists t_import_assignment", + `create table t_import_assignment ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + a_lower varchar(32) collate utf8mb4_general_ci, + a_upper varchar(32) collate utf8mb4_general_ci, + a_concat varchar(80) collate utf8mb4_general_ci, + a_substr varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_lower (a_lower), + key idx_upper (a_upper), + key idx_concat (a_concat), + key idx_substr (a_substr) + )`, + }, + fileName: "t_import_assignment.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: `import into t_import_assignment(@1,@2,@3) + set id=@2, + raw=@3, + a_lower=lower(@3), + a_upper=upper(@3), + a_concat=concat(@2, ':', @3), + a_substr=substr(@3, 1, 2) + from '%s'`, + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + postDMLSQL: []string{ + `insert into t_import_assignment + values ('ddd', 'dbc', lower('dbc'), upper('dbc'), concat('ddd', ':', 'dbc'), substr('dbc', 1, 2))`, + `update t_import_assignment + set raw = 'updated', + a_lower = lower('updated'), + a_upper = upper('updated'), + a_concat = concat(id, ':', 'updated'), + a_substr = substr('updated', 1, 2) + where id = 'ddd'`, + "delete from t_import_assignment where id = 'ddd'", + }, + }, + { + name: "control int primary key and varchar secondary index", + table: "ok_int_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists ok_int_pk_varchar_idx", + `create table ok_int_pk_varchar_idx ( + id int, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_int_pk_varchar_idx.csv", + fileData: "1,abc\n2,bbc\n3,cbc\n", + importSQL: "import into ok_int_pk_varchar_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_int_pk_varchar_idx values (4, 'dbc')", + "update ok_int_pk_varchar_idx set fk = 'updated' where id = 4", + "delete from ok_int_pk_varchar_idx where id = 4", + }, + }, + { + name: "control varchar primary key without secondary index", + table: "ok_varchar_pk_no_secondary", + setupSQL: []string{ + "drop table if exists ok_varchar_pk_no_secondary", + `create table ok_varchar_pk_no_secondary ( + id varchar(32) collate utf8mb4_general_ci, + v int, + primary key (id) clustered + )`, + }, + fileName: "ok_varchar_pk_no_secondary.csv", + fileData: "aaa,1\nbbb,2\nccc,3\n", + importSQL: "import into ok_varchar_pk_no_secondary(id,v) from '%s'", + postDMLSQL: []string{ + "insert into ok_varchar_pk_no_secondary values ('ddd', 4)", + "update ok_varchar_pk_no_secondary set v = 5 where id = 'ddd'", + "delete from ok_varchar_pk_no_secondary where id = 'ddd'", + }, + }, + { + name: "control nonclustered varchar primary key and varchar secondary index", + table: "ok_nonclustered_varchar_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists ok_nonclustered_varchar_pk_varchar_idx", + `create table ok_nonclustered_varchar_pk_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) nonclustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_nonclustered_varchar_pk_varchar_idx.csv", + fileData: "aaa,abc\nbbb,bbc\nccc,cbc\n", + importSQL: "import into ok_nonclustered_varchar_pk_varchar_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_nonclustered_varchar_pk_varchar_idx values ('ddd', 'dbc')", + "update ok_nonclustered_varchar_pk_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from ok_nonclustered_varchar_pk_varchar_idx where id = 'ddd'", + }, + }, + { + name: "control char primary key and char secondary index", + table: "ok_char_pk_char_idx", + setupSQL: []string{ + "drop table if exists ok_char_pk_char_idx", + `create table ok_char_pk_char_idx ( + id char(32) collate utf8mb4_general_ci, + fk char(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_char_pk_char_idx.csv", + fileData: "aaa,abc\nbbb,bbc\nccc,cbc\n", + importSQL: "import into ok_char_pk_char_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_char_pk_char_idx values ('ddd', 'dbc')", + "update ok_char_pk_char_idx set fk = 'updated' where id = 'ddd'", + "delete from ok_char_pk_char_idx where id = 'ddd'", + }, + }, + { + name: "control composite int primary key and int secondary index", + table: "ok_composite_int_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_int_int_pk_int_idx", + `create table ok_composite_int_int_pk_int_idx ( + id1 int, + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_int_int_pk_int_idx.csv", + fileData: "1,10,100\n2,20,200\n3,30,300\n", + importSQL: "import into ok_composite_int_int_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_int_int_pk_int_idx values (4, 40, 400)", + "update ok_composite_int_int_pk_int_idx set fk = 401 where id1 = 4 and id2 = 40", + "delete from ok_composite_int_int_pk_int_idx where id1 = 4 and id2 = 40", + }, + }, + { + name: "control composite varbinary int primary key and int secondary index", + table: "ok_composite_varbinary_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_varbinary_int_pk_int_idx", + `create table ok_composite_varbinary_int_pk_int_idx ( + id1 varbinary(32), + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_varbinary_int_pk_int_idx.csv", + fileData: "aaa,1,10\nbbb,2,20\nccc,3,30\n", + importSQL: "import into ok_composite_varbinary_int_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_varbinary_int_pk_int_idx values ('ddd', 4, 40)", + "update ok_composite_varbinary_int_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 4", + "delete from ok_composite_varbinary_int_pk_int_idx where id1 = 'ddd' and id2 = 4", + }, + }, + { + name: "control composite char primary key and int secondary index", + table: "ok_composite_char_char_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_char_char_pk_int_idx", + `create table ok_composite_char_char_pk_int_idx ( + id1 char(32) collate utf8mb4_general_ci, + id2 char(32) collate utf8mb4_general_ci, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_char_char_pk_int_idx.csv", + fileData: "aaa,ax,10\nbbb,bx,20\nccc,cx,30\n", + importSQL: "import into ok_composite_char_char_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_char_char_pk_int_idx values ('ddd', 'dx', 40)", + "update ok_composite_char_char_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 'dx'", + "delete from ok_composite_char_char_pk_int_idx where id1 = 'ddd' and id2 = 'dx'", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + collate.SetNewCollationEnabledForTest(false) + for _, sql := range tc.setupSQL { + userTK.MustExec(sql) + } + require.NoError(t, objStore.WriteFile(ctx, tc.fileName, []byte(tc.fileData))) + beforeSubmit := taskSubmitCnt.Load() + before := taskRefreshCnt.Load() + fileURL := fmt.Sprintf("s3://next-gen-test/collate-data/%s?%s", tc.fileName, s3Args) + result := userTK.MustQuery(fmt.Sprintf(tc.importSQL, fileURL)).Rows() + require.Len(t, result, 1) + require.Greater(t, taskSubmitCnt.Load(), beforeSubmit) + require.Greater(t, taskRefreshCnt.Load(), before) + + collate.SetNewCollationEnabledForTest(false) + checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") + for _, sql := range tc.postDMLSQL { + userTK.MustExec(sql) + } + checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") + }) + } +} + +func checkImportTableAndIndexes(tk *testkit.TestKit, tableName string, indexes []string, expectedCount string) { + tk.MustExec("admin check table " + tableName) + tk.MustQuery("select count(*) from " + tableName).Check(testkit.Rows(expectedCount)) + for _, indexName := range indexes { + tk.MustQuery(fmt.Sprintf("select count(*) from %s force index(%s)", tableName, indexName)). + Check(testkit.Rows(expectedCount)) + } +} diff --git a/tests/realtikvtest/testkit.go b/tests/realtikvtest/testkit.go index a764f9a98a0fd..2842a65a8fe4a 100644 --- a/tests/realtikvtest/testkit.go +++ b/tests/realtikvtest/testkit.go @@ -106,6 +106,9 @@ func RunTestMain(m *testing.M) { type realtikvStoreOption struct { retainData bool keyspace string + // newCollationsEnabledOnFirstBootstrap is nil unless the test wants to + // bootstrap a keyspace with a specific persisted new-collation setting. + newCollationsEnabledOnFirstBootstrap *bool // only used when keyspace is not SYSTEM, in that case, the SYSTEM store will // be closed together with its domain, if we close it before domain of SYSTEM, // some routine might report errors, and we don't want close twice as the storage @@ -136,6 +139,14 @@ func WithKeyspaceName(name string) RealTiKVStoreOption { } } +// WithNewCollationsEnabledOnFirstBootstrap bootstraps a real TiKV test store +// with the requested persisted new-collation setting. +func WithNewCollationsEnabledOnFirstBootstrap(enabled bool) RealTiKVStoreOption { + return func(opt *realtikvStoreOption) { + opt.newCollationsEnabledOnFirstBootstrap = &enabled + } +} + // WithKeepSystemStore allows the store to keep the SYSTEM keyspace store func WithKeepSystemStore(keep bool) RealTiKVStoreOption { return func(opt *realtikvStoreOption) { @@ -165,6 +176,16 @@ type KSRuntime struct { // PrepareForCrossKSTest prepares the environment for cross keyspace tests. func PrepareForCrossKSTest(t *testing.T, userKSs ...string) map[string]*KSRuntime { + return PrepareForCrossKSTestWithNewCollation(t, nil, userKSs...) +} + +// PrepareForCrossKSTestWithNewCollation prepares cross-keyspace runtimes with +// optional per-keyspace persisted new-collation settings. +func PrepareForCrossKSTestWithNewCollation( + t *testing.T, + newCollationEnabled map[string]bool, + userKSs ...string, +) map[string]*KSRuntime { if !kerneltype.IsNextGen() { t.Fail() } @@ -179,8 +200,18 @@ func PrepareForCrossKSTest(t *testing.T, userKSs ...string) map[string]*KSRuntim ksList := append([]string{keyspace.System}, userKSs...) for _, ks := range ksList { - store, dom := CreateMockStoreAndDomainAndSetup(t, WithKeyspaceName(ks), - WithKeepSystemStore(true), WithKeepSelfStore(true), WithAllocPort(true)) + opts := []RealTiKVStoreOption{ + WithKeyspaceName(ks), + WithKeepSystemStore(true), + WithKeepSelfStore(true), + WithAllocPort(true), + } + if newCollationEnabled != nil { + if enabled, ok := newCollationEnabled[ks]; ok { + opts = append(opts, WithNewCollationsEnabledOnFirstBootstrap(enabled)) + } + } + store, dom := CreateMockStoreAndDomainAndSetup(t, opts...) res[ks] = &KSRuntime{ Store: store, Dom: dom, @@ -239,6 +270,9 @@ func CreateMockStoreAndDomainAndSetup(t *testing.T, opts ...RealTiKVStoreOption) conf.TxnLocalLatches.Enabled = false conf.KeyspaceName = ks conf.Store = config.StoreTypeTiKV + if option.newCollationsEnabledOnFirstBootstrap != nil { + conf.NewCollationsEnabledOnFirstBootstrap = *option.newCollationsEnabledOnFirstBootstrap + } if option.allocPort { conf.Port = uint(mockPortAlloc.Add(1)) }