diff --git a/internal/locate/pd_codec.go b/internal/locate/pd_codec.go index 88e187e07..d155193fc 100644 --- a/internal/locate/pd_codec.go +++ b/internal/locate/pd_codec.go @@ -72,6 +72,19 @@ func NewCodecPDClientWithKeyspace(mode apicodec.Mode, client pd.Client, keyspace return &CodecPDClient{client, codec}, nil } +// NewCodecPDClientWithKeyspaceMeta creates a CodecPDClient in API v2 with keyspace meta. +func NewCodecPDClientWithKeyspaceMeta(mode apicodec.Mode, client pd.Client, keyspaceMeta *keyspacepb.KeyspaceMeta) (*CodecPDClient, error) { + if keyspaceMeta == nil { + return NewCodecPDClient(mode, client), nil + } + + codec, err := apicodec.NewCodecV2(mode, keyspaceMeta) + if err != nil { + return nil, err + } + return &CodecPDClient{client, codec}, nil +} + // GetKeyspaceID attempts to retrieve keyspace ID corresponding to the given keyspace name from PD. func GetKeyspaceID(client pd.Client, name string) (uint32, error) { meta, err := client.LoadKeyspace(context.Background(), apicodec.BuildKeyspaceName(name)) diff --git a/tikv/region.go b/tikv/region.go index 6b5e4874d..fea463865 100644 --- a/tikv/region.go +++ b/tikv/region.go @@ -98,6 +98,9 @@ var NewCodecPDClient = locate.NewCodecPDClient // NewCodecPDClientWithKeyspace creates a CodecPDClient in API v2 with keyspace name. var NewCodecPDClientWithKeyspace = locate.NewCodecPDClientWithKeyspace +// NewCodecPDClientWithKeyspaceMeta creates a CodecPDClient in API v2 with keyspace meta. +var NewCodecPDClientWithKeyspaceMeta = locate.NewCodecPDClientWithKeyspaceMeta + // NewCodecV1 is a constructor for v1 Codec. var NewCodecV1 = apicodec.NewCodecV1