Add Keyonly option for rawKv Scan (#47)

* Add Keyonly option for rawKv Scan.

Signed-off-by: zhongyang.wu <zhongyang.wu@outlook.com>

go.mod

@@ -51,4 +51,4 @@ require (
 	google.golang.org/grpc v1.19.0
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6 // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
-)
+)

mockstore/mocktikv/rpc.go

@@ -476,6 +476,16 @@ func (h *rpcHandler) handleKvRawScan(req *kvrpcpb.RawScanRequest) *kvrpcpb.RawSc
 		endKey = req.EndKey
 	}
 	pairs := rawKV.RawScan(req.GetStartKey(), endKey, int(req.GetLimit()))
+	if req.KeyOnly {
+		//filter values when the client set key only to true.
+		for i := range pairs {
+			pairs[i] = Pair{
+				Key:   pairs[i].Key,
+				Value: nil,
+				Err:   nil,
+			}
+		}
+	}
 	return &kvrpcpb.RawScanResponse{
 		Kvs: convertToPbPairs(pairs),
 	}

rawkv/rawkv.go

@@ -33,6 +33,17 @@ var (
 	ErrMaxScanLimitExceeded = errors.New("limit should be less than MaxRawKVScanLimit")
 )
 
+// ScanOption is used to provide additional information for scaning operaiont
+type ScanOption struct {
+	KeyOnly bool // if true, the result will only contains keys
+}
+
+func DefaultScanOption() ScanOption {
+	return ScanOption{
+		KeyOnly: false,
+	}
+}
+
 // Client is a rawkv client of TiKV server which is used as a key-value storage,
 // only GET/PUT/DELETE commands are supported.
 type Client struct {
@@ -255,10 +266,17 @@ func (c *Client) DeleteRange(ctx context.Context, startKey []byte, endKey []byte
 // If you want to exclude the startKey or include the endKey, append a '\0' to the key. For example, to scan
 // (startKey, endKey], you can write:
 // `Scan(append(startKey, '\0'), append(endKey, '\0'), limit)`.
-func (c *Client) Scan(ctx context.Context, startKey, endKey []byte, limit int) (keys [][]byte, values [][]byte, err error) {
+func (c *Client) Scan(ctx context.Context, startKey, endKey []byte, limit int, options ...ScanOption) (keys [][]byte, values [][]byte, err error) {
 	start := time.Now()
 	defer func() { metrics.RawkvCmdHistogram.WithLabelValues("scan").Observe(time.Since(start).Seconds()) }()
 
+	var option ScanOption
+	if options == nil || len(options) == 0 {
+		option = DefaultScanOption()
+	} else {
+		option = options[0]
+	}
+
 	if limit > c.conf.Raw.MaxScanLimit {
 		return nil, nil, errors.WithStack(ErrMaxScanLimitExceeded)
 	}
@@ -270,6 +288,7 @@ func (c *Client) Scan(ctx context.Context, startKey, endKey []byte, limit int) (
 				StartKey: startKey,
 				EndKey:   endKey,
 				Limit:    uint32(limit - len(keys)),
+				KeyOnly:  option.KeyOnly,
 			},
 		}
 		resp, loc, err := c.sendReq(ctx, startKey, req)
@@ -299,10 +318,17 @@ func (c *Client) Scan(ctx context.Context, startKey, endKey []byte, limit int) (
 // (endKey, startKey], you can write:
 // `ReverseScan(append(startKey, '\0'), append(endKey, '\0'), limit)`.
 // It doesn't support Scanning from "", because locating the last Region is not yet implemented.
-func (c *Client) ReverseScan(ctx context.Context, startKey, endKey []byte, limit int) (keys [][]byte, values [][]byte, err error) {
+func (c *Client) ReverseScan(ctx context.Context, startKey, endKey []byte, limit int, options ...ScanOption) (keys [][]byte, values [][]byte, err error) {
 	start := time.Now()
 	defer func() { metrics.RawkvCmdHistogram.WithLabelValues("reverse_scan").Observe(time.Since(start).Seconds()) }()
 
+	var option ScanOption
+	if options == nil || len(options) == 0 {
+		option = DefaultScanOption()
+	} else {
+		option = options[0]
+	}
+
 	if limit > c.conf.Raw.MaxScanLimit {
 		return nil, nil, errors.WithStack(ErrMaxScanLimitExceeded)
 	}
@@ -315,6 +341,7 @@ func (c *Client) ReverseScan(ctx context.Context, startKey, endKey []byte, limit
 				EndKey:   endKey,
 				Limit:    uint32(limit - len(keys)),
 				Reverse:  true,
+				KeyOnly:  option.KeyOnly,
 			},
 		}
 		resp, loc, err := c.sendReq(ctx, startKey, req)

rawkv/rawkv_test.go

@@ -111,6 +111,18 @@ func (s *testRawKVSuite) mustBatchDelete(c *C, keys [][]byte) {
 	c.Assert(err, IsNil)
 }
 
+func (s *testRawKVSuite) mustScanKeyOnly(c *C, startKey string, limit int, expect ...string) {
+	option := DefaultScanOption()
+	option.KeyOnly = true
+	keys, values, err := s.client.Scan(context.TODO(), []byte(startKey), nil, limit, option)
+	c.Assert(err, IsNil)
+	c.Assert(len(keys), Equals, len(expect))
+	for i := range keys {
+		c.Assert(string(keys[i]), Equals, expect[i])
+		c.Assert(values[i], IsNil)
+	}
+}
+
 func (s *testRawKVSuite) mustScan(c *C, startKey string, limit int, expect ...string) {
 	keys, values, err := s.client.Scan(context.TODO(), []byte(startKey), nil, limit)
 	c.Assert(err, IsNil)
@@ -131,6 +143,18 @@ func (s *testRawKVSuite) mustScanRange(c *C, startKey string, endKey string, lim
 	}
 }
 
+func (s *testRawKVSuite) mustReverseScanKeyOnly(c *C, startKey string, limit int, expect ...string) {
+	option := DefaultScanOption()
+	option.KeyOnly = true
+	keys, values, err := s.client.ReverseScan(context.TODO(), []byte(startKey), nil, limit, option)
+	c.Assert(err, IsNil)
+	c.Assert(len(keys), Equals, len(expect))
+	for i := range keys {
+		c.Assert(string(keys[i]), Equals, expect[i])
+		c.Assert(values[i], IsNil)
+	}
+}
+
 func (s *testRawKVSuite) mustReverseScan(c *C, startKey []byte, limit int, expect ...string) {
 	keys, values, err := s.client.ReverseScan(context.TODO(), startKey, nil, limit)
 	c.Assert(err, IsNil)
@@ -261,6 +285,32 @@ func (s *testRawKVSuite) TestScan(c *C) {
 	check()
 }
 
+func (s *testRawKVSuite) TestScanWithKeyOnly(c *C) {
+	s.mustPut(c, []byte("k1"), []byte("v1"))
+	s.mustPut(c, []byte("k3"), []byte("v3"))
+	s.mustPut(c, []byte("k5"), []byte("v5"))
+	s.mustPut(c, []byte("k7"), []byte("v7"))
+
+	check := func() {
+		s.mustScanKeyOnly(c, "", 1, "k1")
+		s.mustScanKeyOnly(c, "k1", 2, "k1", "k3")
+		s.mustScanKeyOnly(c, "", 10, "k1", "k3", "k5", "k7")
+		s.mustScanKeyOnly(c, "k2", 2, "k3", "k5")
+		s.mustScanKeyOnly(c, "k2", 3, "k3", "k5", "k7")
+	}
+
+	check()
+
+	err := s.split(c, "k", "k2")
+	c.Assert(err, IsNil)
+	check()
+
+	err = s.split(c, "k2", "k5")
+	c.Assert(err, IsNil)
+	check()
+
+}
+
 func (s *testRawKVSuite) TestDeleteRange(c *C) {
 	// Init data
 	testData := map[string]string{}