From 6496ae005db051e2978aad65b9ef65843e4659b1 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Tue, 21 Jun 2016 21:09:49 -0700
Subject: [PATCH] clientv3: add withCount support

---
 clientv3/integration/kv_test.go |  8 ++++++++
 clientv3/kv.go                  |  1 +
 clientv3/op.go                  | 22 +++++++++++++++++++++-
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/clientv3/integration/kv_test.go b/clientv3/integration/kv_test.go
index 97013a516..4dbfd129f 100644
--- a/clientv3/integration/kv_test.go
+++ b/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",
diff --git a/clientv3/kv.go b/clientv3/kv.go
index eca6267de..583c880cc 100644
--- a/clientv3/kv.go
+++ b/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)
diff --git a/clientv3/op.go b/clientv3/op.go
index 2a2e55079..ead1fbea1 100644
--- a/clientv3/op.go
+++ b/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) }