From 94fe2d4d1b9863dd9ccc3131e1655b938cefa87f Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Mon, 25 Mar 2024 23:03:12 -0400
Subject: [PATCH] kv: handle DeleteRange requests in ResponseKeyIterate

This commit updates ResponseKeyIterate to handle DeleteRange requests
that have their ReturnKeys flag set to true.

Epic: None
Release note: None
---
 pkg/kv/kvpb/batch.go      |  8 ++++++++
 pkg/kv/kvpb/batch_test.go | 24 ++++++++++++++++++------
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/pkg/kv/kvpb/batch.go b/pkg/kv/kvpb/batch.go
index 0ec8bcf9239f..17da1c58033b 100644
--- a/pkg/kv/kvpb/batch.go
+++ b/pkg/kv/kvpb/batch.go
@@ -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())
 	}
diff --git a/pkg/kv/kvpb/batch_test.go b/pkg/kv/kvpb/batch_test.go
index 3bc45a086838..c1388b3d3456 100644
--- a/pkg/kv/kvpb/batch_test.go
+++ b/pkg/kv/kvpb/batch_test.go
@@ -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{},
@@ -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) {