Repository: pingcap/tidb
File: pkg/dxf/importinto/conflict_resolution_test.go (Go)

Code before:
	"github.com/pingcap/tidb/pkg/executor/importer"
	"github.com/pingcap/tidb/pkg/ingestor/engineapi"
	tidbkv "github.com/pingcap/tidb/pkg/kv"
	"github.com/pingcap/tidb/pkg/lightning/backend/encode"
	"github.com/pingcap/tidb/pkg/lightning/backend/external"
	"github.com/pingcap/tidb/pkg/objstore/storeapi"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/session"
	"github.com/pingcap/tidb/pkg/table"
	"github.com/pingcap/tidb/pkg/tablecodec"
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
	"github.com/pingcap/tidb/pkg/types"
	"github.com/stretchr/testify/require"
	"github.com/tikv/client-go/v2/tikv"
	"go.uber.org/zap"
)

func writeConflictKVFile(t *testing.T, codec tikv.Codec, kvGroup string, objStore storeapi.Storage, kvs []*external.KVPair) *engineapi.ConflictInfo {
	t.Helper()
	ctx := context.Background()
	var summary *external.WriterSummary
	builder := external.NewWriterBuilder()
	if codec != nil {
		builder.SetTiKVCodec(codec)
	}
	w := builder.
		SetOnCloseFunc(func(s *external.WriterSummary) { summary = s }).
		Build(objStore, "/test", kvGroup)
	for _, kv := range kvs {
		require.NoError(t, w.WriteRow(ctx, kv.Key, kv.Value, nil))
	}
	require.NoError(t, w.Close(ctx))
	require.Len(t, summary.MultipleFilesStats, 1)
	return &engineapi.ConflictInfo{
		Count: uint64(len(kvs)),
		Files: []string{summary.MultipleFilesStats[0].Filenames[0][0]},
	}
}

func generateConflictKVFiles(t *testing.T, tempDir string, tbl table.Table, codec tikv.Codec) importinto.KVGroupConflictInfos {
	t.Helper()
	encodeCfg := &encode.EncodingConfig{
		Table:                tbl,
		UseIdentityAutoRowID: true,
	}
	controller := &importer.LoadDataController{
		ASTArgs: &importer.ASTArgs{},
		Plan:    &importer.Plan{},
		Table:   tbl,
	}

Review comment (suggestion):
```suggestion
	builder := external.NewWriterBuilder().SetTiKVCodec(codec)
```

Code after:
	"github.com/pingcap/tidb/pkg/executor/importer"
	"github.com/pingcap/tidb/pkg/ingestor/engineapi"
	tidbkv "github.com/pingcap/tidb/pkg/kv"
	"github.com/pingcap/tidb/pkg/lightning/backend/encode"
	"github.com/pingcap/tidb/pkg/lightning/backend/external"
	"github.com/pingcap/tidb/pkg/objstore/storeapi"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/session"
	"github.com/pingcap/tidb/pkg/table"
	"github.com/pingcap/tidb/pkg/tablecodec"
	"github.com/pingcap/tidb/pkg/testkit"
	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
	"github.com/pingcap/tidb/pkg/types"
	"github.com/stretchr/testify/require"
	"github.com/tikv/client-go/v2/tikv"
	"go.uber.org/zap"
)

func writeConflictKVFile(t *testing.T, codec tikv.Codec, kvGroup string, objStore storeapi.Storage, kvs []*external.KVPair) *engineapi.ConflictInfo {
	t.Helper()
	ctx := context.Background()
	var summary *external.WriterSummary
	w := external.NewWriterBuilder().
		SetTiKVCodec(codec).
		SetOnCloseFunc(func(s *external.WriterSummary) { summary = s }).
		Build(objStore, "/test", kvGroup)
	for _, kv := range kvs {
		require.NoError(t, w.WriteRow(ctx, kv.Key, kv.Value, nil))
	}
	require.NoError(t, w.Close(ctx))
	require.Len(t, summary.MultipleFilesStats, 1)
	return &engineapi.ConflictInfo{
		Count: uint64(len(kvs)),
		Files: []string{summary.MultipleFilesStats[0].Filenames[0][0]},
	}
}

func generateConflictKVFiles(t *testing.T, tempDir string, tbl table.Table, codec tikv.Codec) importinto.KVGroupConflictInfos {
	t.Helper()
	encodeCfg := &encode.EncodingConfig{
		Table:                tbl,
		UseIdentityAutoRowID: true,
	}
	controller := &importer.LoadDataController{
		ASTArgs: &importer.ASTArgs{},
		Plan:    &importer.Plan{},
		Table:   tbl,
	}
	localEncoder, err := importer.NewTableKVEncoderForDupResolve(encodeCfg, controller)
	require.NoError(t, err)