Repository: tikv/client-go
Pull request #737 (2023-03-10 05:52:09)

attach keyspace_id for manual compact request
Signed-off-by: iosmanthus <myosmanthustree@gmail.com>


Patch:
diff --git a/go.mod b/go.mod
index c9111cf..d2d43d4 100644
--- a/go.mod
+++ b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6
+	github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.14.0
diff --git a/go.sum b/go.sum
index dc21500..25e37eb 100644
--- a/go.sum
+++ b/go.sum
@@ -147,8 +147,8 @@ github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c h1:CgbKAHto5CQgW
 github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6 h1:bgLRG7gPJCq6aduA65ZV7xWQBThTcuarBB9VdfAzV4g=
-github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6/go.mod h1:KUrW1FGoznGMMTssYBu0czfAhn6vQcIrHyZoSC6T990=
+github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924 h1:z6WwBPP0Txmal5zf+H7vf/lSmKZtSS8BTNwiLjEjdnA=
+github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924/go.mod h1:KUrW1FGoznGMMTssYBu0czfAhn6vQcIrHyZoSC6T990=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/internal/apicodec/codec.go b/internal/apicodec/codec.go
index 60b2005..5e3b753 100644
--- a/internal/apicodec/codec.go
+++ b/internal/apicodec/codec.go
@@ -99,9 +99,14 @@ func attachAPICtx(c Codec, req *tikvrpc.Request) (*tikvrpc.Request, error) {
 		mpp := *r.DispatchMPPTask()
 		mpp.Meta.KeyspaceId = ctx.KeyspaceId
 		r.Req = &mpp
+	case tikvrpc.CmdCompact:
+		compact := *r.Compact()
+		compact.KeyspaceId = ctx.KeyspaceId
+		compact.ApiVersion = ctx.ApiVersion
+		r.Req = &compact
 	}
 
-	_ = tikvrpc.AttachContext(&r, ctx)
+	tikvrpc.AttachContext(&r, ctx)
 
 	return &r, nil
 }
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 7e7290f..40f33f0 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -689,8 +689,9 @@ type MPPStreamResponse struct {
 	Lease
 }
 
-// AttachContext sets the request context to the request.
-func AttachContext(req *Request, ctx *kvrpcpb.Context) error {
+// AttachContext sets the request context to the request,
+// return false if encounter unknown request type.
+func AttachContext(req *Request, ctx *kvrpcpb.Context) bool {
 	switch req.Type {
 	case CmdGet:
 		req.Get().Context = ctx
@@ -781,9 +782,9 @@ func AttachContext(req *Request, ctx *kvrpcpb.Context) error {
 	case CmdPrepareFlashbackToVersion:
 		req.PrepareFlashbackToVersion().Context = ctx
 	default:
-		return errors.Errorf("invalid request type %v", req.Type)
+		return false
 	}
-	return nil
+	return true
 }
 
 // SetContext set the Context field for the given req to the specified ctx.
@@ -794,7 +795,10 @@ func SetContext(req *Request, region *metapb.Region, peer *metapb.Peer) error {
 		ctx.RegionEpoch = region.RegionEpoch
 	}
 	ctx.Peer = peer
-	return AttachContext(req, ctx)
+	if !AttachContext(req, ctx) {
+		return errors.Errorf("invalid request type %v", req.Type)
+	}
+	return nil
 }
 
 // GenRegionErrorResp returns corresponding Response with specified RegionError


Test patch:
diff --git a/integration_tests/go.mod b/integration_tests/go.mod
index a603a45..7ee8ee2 100644
--- a/integration_tests/go.mod
+++ b/integration_tests/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20221009092201-b66cddb77c32
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
-	github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6
+	github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924
 	github.com/pingcap/tidb v1.1.0-beta.0.20230207083958-f1d450ff7aa4
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.1
diff --git a/integration_tests/go.sum b/integration_tests/go.sum
index fffc8d1..477647c 100644
--- a/integration_tests/go.sum
+++ b/integration_tests/go.sum
@@ -338,8 +338,8 @@ github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c/go.mod h1:4qGtCB
 github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059 h1:Pe2LbxRmbTfAoKJ65bZLmhahmvHm7n9DUxGRQT00208=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6 h1:bgLRG7gPJCq6aduA65ZV7xWQBThTcuarBB9VdfAzV4g=
-github.com/pingcap/kvproto v0.0.0-20230228041042-1e9aca94bab6/go.mod h1:KUrW1FGoznGMMTssYBu0czfAhn6vQcIrHyZoSC6T990=
+github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924 h1:z6WwBPP0Txmal5zf+H7vf/lSmKZtSS8BTNwiLjEjdnA=
+github.com/pingcap/kvproto v0.0.0-20230312142449-01623096c924/go.mod h1:KUrW1FGoznGMMTssYBu0czfAhn6vQcIrHyZoSC6T990=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20221116035753-734d527bc87c h1:crhkw6DD+07Bg1wYhW5Piw+kYNKZqFQqfC2puUf6gMI=
 github.com/pingcap/log v1.1.1-0.20221116035753-734d527bc87c/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index 14edffd..70ad47a 100644
--- a/internal/locate/region_request_test.go
+++ b/internal/locate/region_request_test.go
@@ -488,6 +488,18 @@ func (s *mockTikvGrpcServer) KvPrepareFlashbackToVersion(context.Context, *kvrpc
 	return nil, errors.New("unreachable")
 }
 
+func (s *mockTikvGrpcServer) EstablishDisaggTask(context.Context, *disaggregated.EstablishDisaggTaskRequest) (*disaggregated.EstablishDisaggTaskResponse, error) {
+	return nil, errors.New("unreachable")
+}
+
+func (s *mockTikvGrpcServer) FetchDisaggPages(*disaggregated.FetchDisaggPagesRequest, tikvpb.Tikv_FetchDisaggPagesServer) error {
+	return errors.New("unreachable")
+}
+
+func (s *mockTikvGrpcServer) TabletSnapshot(_ tikvpb.Tikv_TabletSnapshotServer) error {
+	return errors.New("unreachable")
+}
+
 func (s *testRegionRequestToSingleStoreSuite) TestNoReloadRegionForGrpcWhenCtxCanceled() {
 	// prepare a mock tikv grpc server
 	addr := "localhost:56341"