From 291fbe02baf26fb5f7ebf8e043b8e729de760251 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Wed, 2 Sep 2015 18:08:07 -0700
Subject: [PATCH] consul/state: list keys from the kv with a prefix/separator

---
 consul/state/state_store.go      | 52 ++++++++++++++++++++++++++++++++
 consul/state/state_store_test.go | 38 +++++++++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/consul/state/state_store.go b/consul/state/state_store.go
index 5d1a2cda87..6a1e5ce26f 100644
--- a/consul/state/state_store.go
+++ b/consul/state/state_store.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"strings"
 
 	"github.com/hashicorp/consul/consul/structs"
 	"github.com/hashicorp/go-memdb"
@@ -783,6 +784,57 @@ func (s *StateStore) KVSList(prefix string) (uint64, []string, error) {
 	return lindex, keys, nil
 }
 
+// KVSListKeys is used to query the KV store for keys matching the given prefix.
+// An optional separator may be specified, which can be used to slice off a part
+// of the response so that only a subset of the prefix is returned. In this
+// mode, the keys which are omitted are still counted in the returned index.
+func (s *StateStore) KVSListKeys(prefix, sep string) (uint64, []string, error) {
+	tx := s.db.Txn(false)
+	defer tx.Abort()
+
+	// Fetch keys using the specified prefix
+	entries, err := tx.Get("kvs", "id_prefix", prefix)
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed kvs lookup: %s", err)
+	}
+
+	prefixLen := len(prefix)
+	sepLen := len(sep)
+
+	var keys []string
+	var lindex uint64
+	var last string
+	for entry := entries.Next(); entry != nil; entry = entries.Next() {
+		e := entry.(*structs.DirEntry)
+
+		// Accumulate the high index
+		if e.ModifyIndex > lindex {
+			lindex = e.ModifyIndex
+		}
+
+		// Always accumulate if no separator provided
+		if sepLen == 0 {
+			keys = append(keys, e.Key)
+			continue
+		}
+
+		// Parse and de-duplicate the returned keys based on the
+		// key separator, if provided.
+		after := e.Key[prefixLen:]
+		sepIdx := strings.Index(after, sep)
+		if sepIdx > -1 {
+			key := e.Key[:prefixLen+sepIdx+sepLen]
+			if key != last {
+				keys = append(keys, key)
+				last = key
+			}
+		} else {
+			keys = append(keys, e.Key)
+		}
+	}
+	return lindex, keys, nil
+}
+
 // KVSDelete is used to perform a shallow delete on a single key in the
 // the state store.
 func (s *StateStore) KVSDelete(idx uint64, key string) error {
diff --git a/consul/state/state_store_test.go b/consul/state/state_store_test.go
index da5aa67f0a..d39b09efc0 100644
--- a/consul/state/state_store_test.go
+++ b/consul/state/state_store_test.go
@@ -940,6 +940,44 @@ func TestStateStore_KVSList(t *testing.T) {
 	}
 }
 
+func TestStateStore_KVSListKeys(t *testing.T) {
+	s := testStateStore(t)
+
+	// Create some keys
+	testSetKey(t, s, 1, "foo", "foo")
+	testSetKey(t, s, 2, "foo/bar", "bar")
+	testSetKey(t, s, 3, "foo/bar/baz", "baz")
+	testSetKey(t, s, 4, "foo/bar/zip", "zip")
+	testSetKey(t, s, 5, "foo/bar/zip/zam", "zam")
+	testSetKey(t, s, 6, "foo/bar/zip/zorp", "zorp")
+
+	// Query using a prefix and pass a separator
+	idx, keys, err := s.KVSListKeys("foo/bar/", "/")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if idx != 6 {
+		t.Fatalf("bad index: %d", idx)
+	}
+
+	// Subset of the keys was returned
+	expect := []string{"foo/bar/baz", "foo/bar/zip", "foo/bar/zip/"}
+	if !reflect.DeepEqual(keys, expect) {
+		t.Fatalf("bad keys: %#v", keys)
+	}
+
+	// Listing keys with no separator returns everything.
+	idx, keys, err = s.KVSListKeys("foo", "")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	expect = []string{"foo", "foo/bar", "foo/bar/baz", "foo/bar/zip",
+		"foo/bar/zip/zam", "foo/bar/zip/zorp"}
+	if !reflect.DeepEqual(keys, expect) {
+		t.Fatalf("bad keys: %#v", keys)
+	}
+}
+
 func TestStateStore_KVSDeleteCAS(t *testing.T) {
 	s := testStateStore(t)