Skip to content

*: update tipb and support ResourceGroupTag with keyspace info #61096

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 20, 2025
Merged

*: update tipb and support ResourceGroupTag with keyspace info #61096

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f3931f7
*: support keyspace ID for ResourceGroupTag
zimulala May 13, 2025
52d701f
*: update tipb
zimulala May 13, 2025
4cabe52
*: update bazel
zimulala May 13, 2025
bfd448d
*: tiny update
zimulala May 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions DEPS.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6075,13 +6075,13 @@ def go_deps():
name = "com_github_pingcap_tipb",
build_file_proto_mode = "disable_global",
importpath = "github.com/pingcap/tipb",
sha256 = "5bffa36fa7e42bca9b6de7ba2536f64f5482be1104c838dda78b9beff550942d",
strip_prefix = "github.com/pingcap/[email protected]20250401143359-775c2379cbc7",
sha256 = "4f13b6fc9e5adf6b065369bc9b1a535b4f3e98764002cd2884fe6656c1bfc4e8",
strip_prefix = "github.com/pingcap/[email protected]20250513092957-b555ca3fc078",
urls = [
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250401143359-775c2379cbc7.zip",
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250401143359-775c2379cbc7.zip",
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250401143359-775c2379cbc7.zip",
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250401143359-775c2379cbc7.zip",
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250513092957-b555ca3fc078.zip",
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250513092957-b555ca3fc078.zip",
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250513092957-b555ca3fc078.zip",
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20250513092957-b555ca3fc078.zip",
],
)
go_repository(
Expand Down
2 changes: 1 addition & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ require (
github.com/pingcap/log v1.1.1-0.20250424032633-85a82d016f84
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
github.com/pingcap/tipb v0.0.0-20250401143359-775c2379cbc7
github.com/pingcap/tipb v0.0.0-20250513092957-b555ca3fc078
github.com/prometheus/client_golang v1.22.0
github.com/prometheus/client_model v0.6.2
github.com/prometheus/common v0.63.0
Expand Down
6 changes: 2 additions & 4 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,7 +355,6 @@ github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw
github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
github.com/golang/snappy v0.0.0-20180518054509-2e65f85255db/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
Expand Down Expand Up @@ -691,8 +690,8 @@ github.com/pingcap/log v1.1.1-0.20250424032633-85a82d016f84 h1:ljnSbUq7LOkUtLtDX
github.com/pingcap/log v1.1.1-0.20250424032633-85a82d016f84/go.mod h1:ORfBOFp1eteu2odzsyaxI+b8TzJwgjwyQcGhI+9SfEA=
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE=
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530=
github.com/pingcap/tipb v0.0.0-20250401143359-775c2379cbc7 h1:iRRoPMrpj4jkIzhMT+yabc98bhGeUzAFSjh5HUOn5Yg=
github.com/pingcap/tipb v0.0.0-20250401143359-775c2379cbc7/go.mod h1:zrnYy8vReNODg8G0OiYaX9OK+kpq+rK1jHmvd1DnIWw=
github.com/pingcap/tipb v0.0.0-20250513092957-b555ca3fc078 h1:vkg2gG9XYbgqaxpiI9Qf0w6V4A9sUS+djrpwvLTx+WU=
github.com/pingcap/tipb v0.0.0-20250513092957-b555ca3fc078/go.mod h1:RM8iRcMalzOthG2XJxnNBniM4xFGb/lDwHUwqkaVzt4=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
Expand Down Expand Up @@ -1378,7 +1377,6 @@ google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZi
google.golang.org/grpc v1.20.1/go.mod h1:10oTOabMzJvdu6/UiuZezV6QK5dSlG84ov/aaiqXj38=
google.golang.org/grpc v1.21.1/go.mod h1:oYelfM1adQP15Ek0mdvEgi9Df8B9CZIaU1084ijfRaM=
google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg=
google.golang.org/grpc v1.24.0/go.mod h1:XDChyiUovWa60DnaeDeZmSW86xtLtjtZbwvSiRnRtcA=
google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY=
google.golang.org/grpc v1.26.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
Expand Down
3 changes: 2 additions & 1 deletion pkg/ddl/job_worker.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import (
"github.com/pingcap/tidb/pkg/ddl/systable"
"github.com/pingcap/tidb/pkg/ddl/util"
"github.com/pingcap/tidb/pkg/infoschema"
"github.com/pingcap/tidb/pkg/keyspace"
"github.com/pingcap/tidb/pkg/kv"
"github.com/pingcap/tidb/pkg/meta"
"github.com/pingcap/tidb/pkg/meta/autoid"
Expand Down Expand Up @@ -702,7 +703,7 @@ func (w *ReorgContext) getResourceGroupTaggerForTopSQL() *kv.ResourceGroupTagBui
}

digest := w.cacheDigest
return kv.NewResourceGroupTagBuilder().SetSQLDigest(digest)
return kv.NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).SetSQLDigest(digest)
}

func (w *ReorgContext) ddlJobSourceType() string {
Expand Down
2 changes: 1 addition & 1 deletion pkg/distsql/request_builder_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -884,7 +884,7 @@ func TestRequestBuilderHandle(t *testing.T) {
handles := []kv.Handle{kv.IntHandle(0), kv.IntHandle(2), kv.IntHandle(3), kv.IntHandle(4),
kv.IntHandle(5), kv.IntHandle(10), kv.IntHandle(11), kv.IntHandle(100)}

resourceTagBuilder := kv.NewResourceGroupTagBuilder()
resourceTagBuilder := kv.NewResourceGroupTagBuilder(nil)
tableID := int64(15)
actual, err := (&RequestBuilder{}).SetTableHandles(tableID, handles).
SetDAGRequest(&tipb.DAGRequest{}).
Expand Down
3 changes: 2 additions & 1 deletion pkg/executor/select.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import (
"github.com/pingcap/tidb/pkg/executor/sortexec"
"github.com/pingcap/tidb/pkg/expression"
"github.com/pingcap/tidb/pkg/infoschema"
"github.com/pingcap/tidb/pkg/keyspace"
"github.com/pingcap/tidb/pkg/kv"
"github.com/pingcap/tidb/pkg/meta/model"
"github.com/pingcap/tidb/pkg/parser/ast"
Expand Down Expand Up @@ -322,7 +323,7 @@ func newLockCtx(sctx sessionctx.Context, lockWaitTime int64, numKeys int) (*tikv
}
_, planDigest := seVars.StmtCtx.GetPlanDigest()

return kv.NewResourceGroupTagBuilder().
return kv.NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).
SetPlanDigest(planDigest).
SetSQLDigest(digest).
EncodeTagWithKey(mutation.Key)
Expand Down
6 changes: 3 additions & 3 deletions pkg/infoschema/test/clustertablestest/cluster_tables_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,10 +167,10 @@ func TestTestDataLockWaits(t *testing.T) {
_, digest1 := parser.NormalizeDigest("select * from test_data_lock_waits for update")
_, digest2 := parser.NormalizeDigest("update test_data_lock_waits set f1=1 where id=2")
s.store.(mockstorage.MockLockWaitSetter).SetMockLockWaits([]*deadlock.WaitForEntry{
{Txn: 1, WaitForTxn: 2, Key: []byte("key1"), ResourceGroupTag: kv.NewResourceGroupTagBuilder().SetSQLDigest(digest1).EncodeTagWithKey([]byte(""))},
{Txn: 3, WaitForTxn: 4, Key: []byte("key2"), ResourceGroupTag: kv.NewResourceGroupTagBuilder().SetSQLDigest(digest2).EncodeTagWithKey([]byte(""))},
{Txn: 1, WaitForTxn: 2, Key: []byte("key1"), ResourceGroupTag: kv.NewResourceGroupTagBuilder(nil).SetSQLDigest(digest1).EncodeTagWithKey([]byte(""))},
{Txn: 3, WaitForTxn: 4, Key: []byte("key2"), ResourceGroupTag: kv.NewResourceGroupTagBuilder(nil).SetSQLDigest(digest2).EncodeTagWithKey([]byte(""))},
// Invalid digests
{Txn: 5, WaitForTxn: 6, Key: []byte("key3"), ResourceGroupTag: kv.NewResourceGroupTagBuilder().EncodeTagWithKey([]byte(""))},
{Txn: 5, WaitForTxn: 6, Key: []byte("key3"), ResourceGroupTag: kv.NewResourceGroupTagBuilder(nil).EncodeTagWithKey([]byte(""))},
{Txn: 7, WaitForTxn: 8, Key: []byte("key4"), ResourceGroupTag: []byte("asdfghjkl")},
})
tk := s.newTestKitWithRoot(t)
Expand Down
3 changes: 3 additions & 0 deletions pkg/keyspace/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ go_library(
visibility = ["//visibility:public"],
deps = [
"//pkg/config",
"//pkg/config/kerneltype",
"@com_github_pingcap_kvproto//pkg/kvrpcpb",
"@com_github_tikv_client_go_v2//tikv",
"@org_uber_go_zap//:zap",
Expand All @@ -20,8 +21,10 @@ go_test(
srcs = ["keyspace_test.go"],
embed = [":keyspace"],
flaky = True,
shard_count = 3,
deps = [
"//pkg/config",
"//pkg/config/kerneltype",
"@com_github_stretchr_testify//require",
],
)
16 changes: 16 additions & 0 deletions pkg/keyspace/keyspace.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,11 @@ package keyspace

import (
"fmt"
"strconv"

"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/pingcap/tidb/pkg/config"
"github.com/pingcap/tidb/pkg/config/kerneltype"
"github.com/tikv/client-go/v2/tikv"
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
Expand Down Expand Up @@ -54,6 +56,20 @@ func GetKeyspaceNameBySettings() (keyspaceName string) {
return keyspaceName
}

// GetKeyspaceIDBySettings is used to get Keyspace ID setting.
func GetKeyspaceIDBySettings() (keyspaceID *uint32) {
keyspaceName := config.GetGlobalKeyspaceName()
if !IsKeyspaceNameEmpty(keyspaceName) && kerneltype.IsNextGen() {
keyspaceIDU64, err := strconv.ParseUint(keyspaceName, 10, 32)
if err != nil {
return nil
}
keyspaceIDU32 := uint32(keyspaceIDU64)
return &keyspaceIDU32
}
return nil
}

// IsKeyspaceNameEmpty is used to determine whether keyspaceName is set.
func IsKeyspaceNameEmpty(keyspaceName string) bool {
return keyspaceName == ""
Expand Down
19 changes: 19 additions & 0 deletions pkg/keyspace/keyspace_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"testing"

"github.com/pingcap/tidb/pkg/config"
"github.com/pingcap/tidb/pkg/config/kerneltype"
"github.com/stretchr/testify/require"
)

Expand Down Expand Up @@ -50,3 +51,21 @@ func TestNoKeyspaceNameSet(t *testing.T) {
require.Equal(t, "", getKeyspaceName)
require.Equal(t, true, IsKeyspaceNameEmpty(getKeyspaceName))
}

func TestNoKeyspaceIDBySettings(t *testing.T) {
config.UpdateGlobal(func(conf *config.Config) {
conf.KeyspaceName = ""
})
getKeyspaceID := GetKeyspaceIDBySettings()
require.Nil(t, getKeyspaceID)

config.UpdateGlobal(func(conf *config.Config) {
conf.KeyspaceName = "123"
})
getKeyspaceID = GetKeyspaceIDBySettings()
if kerneltype.IsNextGen() {
require.Equal(t, uint32(123), *getKeyspaceID)
} else {
require.Nil(t, getKeyspaceID)
}
}
1 change: 1 addition & 0 deletions pkg/kv/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ go_test(
flaky = True,
shard_count = 24,
deps = [
"//pkg/keyspace",
"//pkg/meta/model",
"//pkg/parser",
"//pkg/parser/mysql",
Expand Down
7 changes: 4 additions & 3 deletions pkg/kv/kv.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -785,12 +785,13 @@ const (
type ResourceGroupTagBuilder struct {
sqlDigest *parser.Digest
planDigest *parser.Digest
keyspaceID *uint32
accessKey []byte
}

// NewResourceGroupTagBuilder creates a new ResourceGroupTagBuilder.
func NewResourceGroupTagBuilder() *ResourceGroupTagBuilder {
return &ResourceGroupTagBuilder{}
func NewResourceGroupTagBuilder(keyspaceID *uint32) *ResourceGroupTagBuilder {
return &ResourceGroupTagBuilder{keyspaceID: keyspaceID}
}

// SetSQLDigest sets the sql digest for the request.
Expand All @@ -814,7 +815,7 @@ func (b *ResourceGroupTagBuilder) BuildProtoTagger() tikvrpc.ResourceGroupTagger

// EncodeTagWithKey encodes the resource group tag, returns the encoded bytes.
func (b *ResourceGroupTagBuilder) EncodeTagWithKey(key []byte) []byte {
tag := &tipb.ResourceGroupTag{}
tag := &tipb.ResourceGroupTag{KeyspaceId: b.keyspaceID}
if b.sqlDigest != nil {
tag.SqlDigest = b.sqlDigest.Bytes()
}
Expand Down
24 changes: 20 additions & 4 deletions pkg/kv/kv_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,10 @@ import (
"math/rand"
"testing"

"github.com/pingcap/tidb/pkg/keyspace"
"github.com/pingcap/tidb/pkg/parser"
"github.com/pingcap/tidb/pkg/util/resourcegrouptag"
"github.com/pingcap/tipb/go-tipb"
"github.com/stretchr/testify/require"
)

Expand All @@ -34,31 +36,45 @@ func genRandHex(length int) []byte {

func TestResourceGroupTagEncoding(t *testing.T) {
sqlDigest := parser.NewDigest(nil)
tag := NewResourceGroupTagBuilder().SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
tag := NewResourceGroupTagBuilder(nil).SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
require.Len(t, tag, 2)

decodedSQLDigest, err := resourcegrouptag.DecodeResourceGroupTag(tag)
require.NoError(t, err)
require.Len(t, decodedSQLDigest, 0)
resTag := &tipb.ResourceGroupTag{}
err = resTag.Unmarshal(tag)
require.NoError(t, err)
require.Nil(t, resTag.KeyspaceId)

sqlDigest = parser.NewDigest([]byte{'a', 'a'})
tag = NewResourceGroupTagBuilder().SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
tag = NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
// version(1) + prefix(1) + length(1) + content(2hex -> 1byte)
require.Len(t, tag, 6)

decodedSQLDigest, err = resourcegrouptag.DecodeResourceGroupTag(tag)
require.NoError(t, err)
require.Equal(t, sqlDigest.Bytes(), decodedSQLDigest)

keyspaceID := uint32(123)
sqlDigest = parser.NewDigest(genRandHex(64))
tag = NewResourceGroupTagBuilder().SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
tag = NewResourceGroupTagBuilder(&keyspaceID).SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
decodedSQLDigest, err = resourcegrouptag.DecodeResourceGroupTag(tag)
require.NoError(t, err)
require.Equal(t, sqlDigest.Bytes(), decodedSQLDigest)
resTag = &tipb.ResourceGroupTag{}
err = resTag.Unmarshal(tag)
require.NoError(t, err)
require.NotNil(t, resTag.KeyspaceId)
require.Equal(t, resTag.KeyspaceId, &keyspaceID)

sqlDigest = parser.NewDigest(genRandHex(510))
tag = NewResourceGroupTagBuilder().SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
tag = NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).SetSQLDigest(sqlDigest).EncodeTagWithKey([]byte(""))
decodedSQLDigest, err = resourcegrouptag.DecodeResourceGroupTag(tag)
require.NoError(t, err)
require.Equal(t, sqlDigest.Bytes(), decodedSQLDigest)
resTag = &tipb.ResourceGroupTag{}
err = resTag.Unmarshal(tag)
require.NoError(t, err)
require.Nil(t, resTag.KeyspaceId)
}
1 change: 1 addition & 0 deletions pkg/server/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ go_library(
"//pkg/expression",
"//pkg/extension",
"//pkg/infoschema",
"//pkg/keyspace",
"//pkg/kv",
"//pkg/metrics",
"//pkg/param",
Expand Down
3 changes: 2 additions & 1 deletion pkg/server/conn.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ import (
"github.com/pingcap/tidb/pkg/executor"
"github.com/pingcap/tidb/pkg/extension"
"github.com/pingcap/tidb/pkg/infoschema"
"github.com/pingcap/tidb/pkg/keyspace"
"github.com/pingcap/tidb/pkg/kv"
"github.com/pingcap/tidb/pkg/metrics"
"github.com/pingcap/tidb/pkg/parser"
Expand Down Expand Up @@ -2003,7 +2004,7 @@ func setResourceGroupTaggerForMultiStmtPrefetch(snapshot kv.Snapshot, sqls strin
normalized, digest := parser.NormalizeDigest(sqls)
topsql.AttachAndRegisterSQLInfo(context.Background(), normalized, digest, false)
if len(normalized) != 0 {
snapshot.SetOption(kv.ResourceGroupTagger, kv.NewResourceGroupTagBuilder().SetSQLDigest(digest))
snapshot.SetOption(kv.ResourceGroupTagger, kv.NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).SetSQLDigest(digest))
}
}

Expand Down
1 change: 1 addition & 0 deletions pkg/sessionctx/stmtctx/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ go_library(
deps = [
"//pkg/distsql/context",
"//pkg/errctx",
"//pkg/keyspace",
"//pkg/kv",
"//pkg/meta/model",
"//pkg/parser",
Expand Down
3 changes: 2 additions & 1 deletion pkg/sessionctx/stmtctx/stmtctx.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import (
"github.com/pingcap/failpoint"
distsqlctx "github.com/pingcap/tidb/pkg/distsql/context"
"github.com/pingcap/tidb/pkg/errctx"
"github.com/pingcap/tidb/pkg/keyspace"
"github.com/pingcap/tidb/pkg/kv"
"github.com/pingcap/tidb/pkg/meta/model"
"github.com/pingcap/tidb/pkg/parser"
Expand Down Expand Up @@ -761,7 +762,7 @@ func (sc *StatementContext) SetBinaryPlan(binaryPlan string) {

// GetResourceGroupTagger returns the implementation of kv.ResourceGroupTagBuilder related to self.
func (sc *StatementContext) GetResourceGroupTagger() *kv.ResourceGroupTagBuilder {
tagger := kv.NewResourceGroupTagBuilder().SetPlanDigest(sc.planDigest)
tagger := kv.NewResourceGroupTagBuilder(keyspace.GetKeyspaceIDBySettings()).SetPlanDigest(sc.planDigest)
normalized, digest := sc.SQLDigest()
if len(normalized) > 0 {
tagger.SetSQLDigest(digest)
Expand Down