diff --git a/etcdctl/ctlv3/command/get_command.go b/etcdctl/ctlv3/command/get_command.go index 753d75ba799..bb7347051f6 100644 --- a/etcdctl/ctlv3/command/get_command.go +++ b/etcdctl/ctlv3/command/get_command.go @@ -31,6 +31,7 @@ var ( getFromKey bool getRev int64 getKeysOnly bool + getCountOnly bool printValueOnly bool ) @@ -50,6 +51,7 @@ func NewGetCommand() *cobra.Command { cmd.Flags().BoolVar(&getFromKey, "from-key", false, "Get keys that are greater than or equal to the given key using byte compare") cmd.Flags().Int64Var(&getRev, "rev", 0, "Specify the kv revision") cmd.Flags().BoolVar(&getKeysOnly, "keys-only", false, "Get only the keys") + cmd.Flags().BoolVar(&getCountOnly, "count-only", false, "Get only the count") cmd.Flags().BoolVar(&printValueOnly, "print-value-only", false, `Only write values when using the "simple" output format`) return cmd } @@ -64,6 +66,12 @@ func getCommandFunc(cmd *cobra.Command, args []string) { ExitWithError(ExitError, err) } + if getCountOnly { + if _, fields := display.(*fieldsPrinter); !fields { + ExitWithError(ExitBadArgs, fmt.Errorf("--count-only is only for `--write-out=fields`")) + } + } + if printValueOnly { dp, simple := (display).(*simplePrinter) if !simple { @@ -83,6 +91,10 @@ func getGetOp(args []string) (string, []clientv3.OpOption) { ExitWithError(ExitBadArgs, fmt.Errorf("`--prefix` and `--from-key` cannot be set at the same time, choose one")) } + if getKeysOnly && getCountOnly { + ExitWithError(ExitBadArgs, fmt.Errorf("`--keys-only` and `--count-only` cannot be set at the same time, choose one")) + } + opts := []clientv3.OpOption{} switch getConsistency { case "s": @@ -159,5 +171,9 @@ func getGetOp(args []string) (string, []clientv3.OpOption) { opts = append(opts, clientv3.WithKeysOnly()) } + if getCountOnly { + opts = append(opts, clientv3.WithCountOnly()) + } + return key, opts } diff --git a/mvcc/index.go b/mvcc/index.go index f8cc6df88cf..f0638523f72 100644 --- a/mvcc/index.go +++ b/mvcc/index.go @@ -26,6 +26,7 @@ type index interface { Get(key []byte, atRev int64) (rev, created revision, ver int64, err error) Range(key, end []byte, atRev int64) ([][]byte, []revision) Revisions(key, end []byte, atRev int64) []revision + CountRevisions(key, end []byte, atRev int64) int Put(key []byte, rev revision) Tombstone(key []byte, rev revision) error RangeSince(key, end []byte, rev int64) []revision @@ -119,6 +120,23 @@ func (ti *treeIndex) Revisions(key, end []byte, atRev int64) (revs []revision) { return revs } +func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int { + if end == nil { + _, _, _, err := ti.Get(key, atRev) + if err != nil { + return 0 + } + return 1 + } + total := 0 + ti.visit(key, end, func(ki *keyIndex) { + if _, _, _, err := ki.get(ti.lg, atRev); err == nil { + total++ + } + }) + return total +} + func (ti *treeIndex) Range(key, end []byte, atRev int64) (keys [][]byte, revs []revision) { if end == nil { rev, _, _, err := ti.Get(key, atRev) diff --git a/mvcc/index_test.go b/mvcc/index_test.go index f9431ae314a..8befe14cda0 100644 --- a/mvcc/index_test.go +++ b/mvcc/index_test.go @@ -194,6 +194,80 @@ func TestIndexRangeSince(t *testing.T) { } } +func TestIndexRevision(t *testing.T) { + allKeys := [][]byte{[]byte("foo"), []byte("foo1"), []byte("foo2"), []byte("foo2"), []byte("foo1"), []byte("foo")} + allRevs := []revision{{main: 1}, {main: 2}, {main: 3}, {main: 4}, {main: 5}, {main: 6}} + + ti := newTreeIndex(zap.NewExample()) + for i := range allKeys { + ti.Put(allKeys[i], allRevs[i]) + } + + tests := []struct { + key, end []byte + atRev int64 + wrevs []revision + wcounts int + }{ + // single key that not found + { + []byte("bar"), nil, 6, nil, 0, + }, + // single key that found + { + []byte("foo"), nil, 6, []revision{{main: 6}}, 1, + }, + // various range keys, fixed atRev + { + []byte("foo"), []byte("foo1"), 6, []revision{{main: 6}}, 1, + }, + { + []byte("foo"), []byte("foo2"), 6, []revision{{main: 6}, {main: 5}}, 2, + }, + { + []byte("foo"), []byte("fop"), 6, []revision{{main: 6}, {main: 5}, {main: 4}}, 3, + }, + { + []byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2, + }, + { + []byte("foo2"), []byte("fop"), 6, []revision{{main: 4}}, 1, + }, + { + []byte("foo3"), []byte("fop"), 6, nil, 0, + }, + // fixed range keys, various atRev + { + []byte("foo1"), []byte("fop"), 1, nil, 0, + }, + { + []byte("foo1"), []byte("fop"), 2, []revision{{main: 2}}, 1, + }, + { + []byte("foo1"), []byte("fop"), 3, []revision{{main: 2}, {main: 3}}, 2, + }, + { + []byte("foo1"), []byte("fop"), 4, []revision{{main: 2}, {main: 4}}, 2, + }, + { + []byte("foo1"), []byte("fop"), 5, []revision{{main: 5}, {main: 4}}, 2, + }, + { + []byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2, + }, + } + for i, tt := range tests { + revs := ti.Revisions(tt.key, tt.end, tt.atRev) + if !reflect.DeepEqual(revs, tt.wrevs) { + t.Errorf("#%d: revs = %+v, want %+v", i, revs, tt.wrevs) + } + count := ti.CountRevisions(tt.key, tt.end, tt.atRev) + if count != tt.wcounts { + t.Errorf("#%d: count = %d, want %v", i, count, tt.wcounts) + } + } +} + func TestIndexCompactAndKeep(t *testing.T) { maxRev := int64(20) tests := []struct { diff --git a/mvcc/kvstore_test.go b/mvcc/kvstore_test.go index eb9b1f130f4..bef9d365bd0 100644 --- a/mvcc/kvstore_test.go +++ b/mvcc/kvstore_test.go @@ -941,6 +941,11 @@ func (i *fakeIndex) Revisions(key, end []byte, atRev int64) []revision { return rev } +func (i *fakeIndex) CountRevisions(key, end []byte, atRev int64) int { + _, rev := i.Range(key, end, atRev) + return len(rev) +} + func (i *fakeIndex) Get(key []byte, atRev int64) (rev, created revision, ver int64, err error) { i.Recorder.Record(testutil.Action{Name: "get", Params: []interface{}{key, atRev}}) r := <-i.indexGetRespc diff --git a/mvcc/kvstore_txn.go b/mvcc/kvstore_txn.go index 5ba7f2bb007..5a363b2d184 100644 --- a/mvcc/kvstore_txn.go +++ b/mvcc/kvstore_txn.go @@ -125,15 +125,16 @@ func (tr *storeTxnRead) rangeKeys(key, end []byte, curRev int64, ro RangeOptions if rev < tr.s.compactMainRev { return &RangeResult{KVs: nil, Count: -1, Rev: 0}, ErrCompacted } - + if ro.Count { + total := tr.s.kvindex.CountRevisions(key, end, rev) + tr.trace.Step("count revisions from in-memory index tree") + return &RangeResult{KVs: nil, Count: total, Rev: curRev}, nil + } revpairs := tr.s.kvindex.Revisions(key, end, rev) tr.trace.Step("range keys from in-memory index tree") if len(revpairs) == 0 { return &RangeResult{KVs: nil, Count: 0, Rev: curRev}, nil } - if ro.Count { - return &RangeResult{KVs: nil, Count: len(revpairs), Rev: curRev}, nil - } limit := int(ro.Limit) if limit <= 0 || limit > len(revpairs) { diff --git a/tests/e2e/ctl_v3_kv_test.go b/tests/e2e/ctl_v3_kv_test.go index e0e3c8c1fdc..3b705c618ea 100644 --- a/tests/e2e/ctl_v3_kv_test.go +++ b/tests/e2e/ctl_v3_kv_test.go @@ -41,9 +41,10 @@ func TestCtlV3GetPeerTLS(t *testing.T) { testCtl(t, getTest, withCfg(confi func TestCtlV3GetTimeout(t *testing.T) { testCtl(t, getTest, withDialTimeout(0)) } func TestCtlV3GetQuorum(t *testing.T) { testCtl(t, getTest, withQuorum()) } -func TestCtlV3GetFormat(t *testing.T) { testCtl(t, getFormatTest) } -func TestCtlV3GetRev(t *testing.T) { testCtl(t, getRevTest) } -func TestCtlV3GetKeysOnly(t *testing.T) { testCtl(t, getKeysOnlyTest) } +func TestCtlV3GetFormat(t *testing.T) { testCtl(t, getFormatTest) } +func TestCtlV3GetRev(t *testing.T) { testCtl(t, getRevTest) } +func TestCtlV3GetKeysOnly(t *testing.T) { testCtl(t, getKeysOnlyTest) } +func TestCtlV3GetCountOnly(t *testing.T) { testCtl(t, getCountOnlyTest) } func TestCtlV3Del(t *testing.T) { testCtl(t, delTest) } func TestCtlV3DelNoTLS(t *testing.T) { testCtl(t, delTest, withCfg(configNoTLS)) } @@ -235,6 +236,44 @@ func getKeysOnlyTest(cx ctlCtx) { } } +func getCountOnlyTest(cx ctlCtx) { + cmdArgs := append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...) + if err := spawnWithExpects(cmdArgs, "\"Count\" : 0"); err != nil { + cx.t.Fatal(err) + } + if err := ctlV3Put(cx, "key", "val", ""); err != nil { + cx.t.Fatal(err) + } + cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...) + if err := spawnWithExpects(cmdArgs, "\"Count\" : 1"); err != nil { + cx.t.Fatal(err) + } + if err := ctlV3Put(cx, "key1", "val", ""); err != nil { + cx.t.Fatal(err) + } + if err := ctlV3Put(cx, "key1", "val", ""); err != nil { + cx.t.Fatal(err) + } + cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...) + if err := spawnWithExpects(cmdArgs, "\"Count\" : 2"); err != nil { + cx.t.Fatal(err) + } + if err := ctlV3Put(cx, "key2", "val", ""); err != nil { + cx.t.Fatal(err) + } + cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...) + if err := spawnWithExpects(cmdArgs, "\"Count\" : 3"); err != nil { + cx.t.Fatal(err) + } + expected := []string{ + "\"Count\" : 3", + } + cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key3", "--prefix", "--write-out=fields"}...) + if err := spawnWithExpects(cmdArgs, expected...); err == nil { + cx.t.Fatal(err) + } +} + func delTest(cx ctlCtx) { tests := []struct { puts []kv