Merge pull request #5745 from xiang90/count_client

clientv3: add withCount support

clientv3/integration/kv_test.go

@@ -201,6 +201,14 @@ func TestKVRange(t *testing.T) {
 				{Key: []byte("a"), Value: nil, CreateRevision: 2, ModRevision: 2, Version: 1},
 			},
 		},
+		// range all with countOnly
+		{
+			"a", "x",
+			2,
+			[]clientv3.OpOption{clientv3.WithCountOnly()},
+
+			nil,
+		},
 		// range all with SortByKey, SortAscend
 		{
 			"a", "x",

clientv3/kv.go

@@ -141,6 +141,7 @@ func (kv *kv) do(ctx context.Context, op Op) (OpResponse, error) {
 			Revision:     op.rev,
 			Serializable: op.serializable,
 			KeysOnly:     op.keysOnly,
+			CountOnly:    op.countOnly,
 		}
 		if op.sort != nil {
 			r.SortOrder = pb.RangeRequest_SortOrder(op.sort.Order)

clientv3/op.go

@@ -42,6 +42,7 @@ type Op struct {
 	sort         *SortOption
 	serializable bool
 	keysOnly     bool
+	countOnly    bool
 
 	// for range, watch
 	rev int64
@@ -57,7 +58,15 @@ type Op struct {
 func (op Op) toRequestOp() *pb.RequestOp {
 	switch op.t {
 	case tRange:
-		r := &pb.RangeRequest{Key: op.key, RangeEnd: op.end, Limit: op.limit, Revision: op.rev, Serializable: op.serializable}
+		r := &pb.RangeRequest{
+			Key:          op.key,
+			RangeEnd:     op.end,
+			Limit:        op.limit,
+			Revision:     op.rev,
+			Serializable: op.serializable,
+			KeysOnly:     op.keysOnly,
+			CountOnly:    op.countOnly,
+		}
 		if op.sort != nil {
 			r.SortOrder = pb.RangeRequest_SortOrder(op.sort.Order)
 			r.SortTarget = pb.RangeRequest_SortTarget(op.sort.Target)
@@ -98,6 +107,8 @@ func OpDelete(key string, opts ...OpOption) Op {
 		panic("unexpected sort in delete")
 	case ret.serializable:
 		panic("unexpected serializable in delete")
+	case ret.countOnly:
+		panic("unexpected countOnly in delete")
 	}
 	return ret
 }
@@ -116,6 +127,8 @@ func OpPut(key, val string, opts ...OpOption) Op {
 		panic("unexpected sort in put")
 	case ret.serializable:
 		panic("unexpected serializable in put")
+	case ret.countOnly:
+		panic("unexpected countOnly in delete")
 	}
 	return ret
 }
@@ -132,6 +145,8 @@ func opWatch(key string, opts ...OpOption) Op {
 		panic("unexpected sort in watch")
 	case ret.serializable:
 		panic("unexpected serializable in watch")
+	case ret.countOnly:
+		panic("unexpected countOnly in delete")
 	}
 	return ret
 }
@@ -215,6 +230,11 @@ func WithKeysOnly() OpOption {
 	return func(op *Op) { op.keysOnly = true }
 }
 
+// WithCountOnly makes the 'Get' request return only the count of keys.
+func WithCountOnly() OpOption {
+	return func(op *Op) { op.countOnly = true }
+}
+
 // WithFirstCreate gets the key with the oldest creation revision in the request range.
 func WithFirstCreate() []OpOption { return withTop(SortByCreateRevision, SortAscend) }