Skip to content

Commit

Permalink
kv: handle DeleteRange requests in ResponseKeyIterate
Browse files Browse the repository at this point in the history
This commit updates ResponseKeyIterate to handle DeleteRange requests
that have their ReturnKeys flag set to true.

Epic: None
Release note: None
  • Loading branch information
nvanbenschoten committed Mar 26, 2024
1 parent b85aaa8 commit 94fe2d4
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 6 deletions.
8 changes: 8 additions & 0 deletions pkg/kv/kvpb/batch.go
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,14 @@ func ResponseKeyIterate(req Request, resp Response, fn func(roachpb.Key)) error
return errors.AssertionFailedf("unexpectedly non-empty ColBatches")
}
}
case *DeleteRangeResponse:
if !req.(*DeleteRangeRequest).ReturnKeys {
return errors.AssertionFailedf("cannot iterate over response keys of DeleteRange " +
"request when ReturnKeys=false")
}
for _, k := range v.Keys {
fn(k)
}
default:
return errors.Errorf("cannot iterate over response keys of %s request", req.Method())
}
Expand Down
24 changes: 18 additions & 6 deletions pkg/kv/kvpb/batch_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -471,6 +471,24 @@ func TestResponseKeyIterate(t *testing.T) {
resp: &ReverseScanResponse{},
expErr: "unexpectedly called ResponseKeyIterate on a ReverseScanRequest with COL_BATCH_RESPONSE scan format",
},
{
name: "delete-range, no response keys",
req: &DeleteRangeRequest{RequestHeader: rangeHeader, ReturnKeys: true},
resp: &DeleteRangeResponse{Keys: nil},
expKeys: nil,
},
{
name: "delete-range, some response keys",
req: &DeleteRangeRequest{RequestHeader: rangeHeader, ReturnKeys: true},
resp: &DeleteRangeResponse{Keys: []roachpb.Key{keyB, keyC}},
expKeys: []roachpb.Key{keyB, keyC},
},
{
name: "delete-range, return keys disabled",
req: &DeleteRangeRequest{RequestHeader: rangeHeader, ReturnKeys: false},
resp: &DeleteRangeResponse{},
expErr: "cannot iterate over response keys of DeleteRange request when ReturnKeys=false",
},
{
name: "put",
req: &PutRequest{},
Expand Down Expand Up @@ -501,12 +519,6 @@ func TestResponseKeyIterate(t *testing.T) {
resp: &DeleteResponse{},
expErr: "cannot iterate over response keys of Delete request",
},
{
name: "delete-range",
req: &DeleteRangeRequest{},
resp: &DeleteRangeResponse{},
expErr: "cannot iterate over response keys of DeleteRange request",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
Expand Down

0 comments on commit 94fe2d4

Please sign in to comment.