From dcbb55b87403ca0fa250db93e82967a16186295c Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 28 Apr 2014 16:26:31 -0700
Subject: [PATCH] consul: Adding a key listing mechanism

---
 consul/state_store.go      |  47 +++++++++++++++++
 consul/state_store_test.go | 100 +++++++++++++++++++++++++++++++++++++
 2 files changed, 147 insertions(+)

diff --git a/consul/state_store.go b/consul/state_store.go
index fa19b7fb05..d229dea407 100644
--- a/consul/state_store.go
+++ b/consul/state_store.go
@@ -9,6 +9,7 @@ import (
 	"log"
 	"os"
 	"runtime"
+	"strings"
 )
 
 const (
@@ -245,6 +246,7 @@ func (s *StateStore) initialize() error {
 		"NodeDump":          MDBTables{s.nodeTable, s.serviceTable, s.checkTable},
 		"KVSGet":            MDBTables{s.kvsTable},
 		"KVSList":           MDBTables{s.kvsTable},
+		"KVSListKeys":       MDBTables{s.kvsTable},
 	}
 	return nil
 }
@@ -902,6 +904,51 @@ func (s *StateStore) KVSList(prefix string) (uint64, structs.DirEntries, error)
 	return idx, ents, err
 }
 
+// KVSListKeys is used to list keys with a prefix, and up to a given seperator
+func (s *StateStore) KVSListKeys(prefix, seperator string) (uint64, []string, error) {
+	tx, err := s.kvsTable.StartTxn(true, nil)
+	if err != nil {
+		return 0, nil, err
+	}
+	defer tx.Abort()
+
+	idx, err := s.kvsTable.LastIndexTxn(tx)
+	if err != nil {
+		return 0, nil, err
+	}
+
+	// Aggregate the stream
+	stream := make(chan interface{}, 128)
+	done := make(chan struct{})
+	var keys []string
+	go func() {
+		prefixLen := len(prefix)
+		last := ""
+		for raw := range stream {
+			ent := raw.(*structs.DirEntry)
+			after := ent.Key[prefixLen:]
+
+			// Check for the seperator
+			if idx := strings.Index(after, seperator); idx >= 0 {
+				toSep := ent.Key[:prefixLen+idx+1]
+				if last != toSep {
+					keys = append(keys, toSep)
+					last = toSep
+				}
+
+			} else {
+				keys = append(keys, ent.Key)
+			}
+		}
+		close(done)
+	}()
+
+	// Start the stream, and wait for completion
+	err = s.kvsTable.StreamTxn(stream, tx, "id_prefix", prefix)
+	<-done
+	return idx, keys, err
+}
+
 // KVSDelete is used to delete a KVS entry
 func (s *StateStore) KVSDelete(index uint64, key string) error {
 	return s.kvsDeleteWithIndex(index, "id", key)
diff --git a/consul/state_store_test.go b/consul/state_store_test.go
index 7c30ede634..f5157771c5 100644
--- a/consul/state_store_test.go
+++ b/consul/state_store_test.go
@@ -1401,6 +1401,106 @@ func TestKVS_List(t *testing.T) {
 	}
 }
 
+func TestKVS_ListKeys(t *testing.T) {
+	store, err := testStateStore()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer store.Close()
+
+	// Should not exist
+	idx, keys, err := store.KVSListKeys("", "/")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if idx != 0 {
+		t.Fatalf("bad: %v", idx)
+	}
+	if len(keys) != 0 {
+		t.Fatalf("bad: %v", keys)
+	}
+
+	// Create the entries
+	d := &structs.DirEntry{Key: "/web/a", Flags: 42, Value: []byte("test")}
+	if err := store.KVSSet(1000, d); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	d = &structs.DirEntry{Key: "/web/b", Flags: 42, Value: []byte("test")}
+	if err := store.KVSSet(1001, d); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	d = &structs.DirEntry{Key: "/web/sub/c", Flags: 42, Value: []byte("test")}
+	if err := store.KVSSet(1002, d); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Should list
+	idx, keys, err = store.KVSListKeys("", "/")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if idx != 1002 {
+		t.Fatalf("bad: %v", idx)
+	}
+	if len(keys) != 1 {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[0] != "/" {
+		t.Fatalf("bad: %v", keys)
+	}
+
+	// Should list just web
+	idx, keys, err = store.KVSListKeys("/", "/")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if idx != 1002 {
+		t.Fatalf("bad: %v", idx)
+	}
+	if len(keys) != 1 {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[0] != "/web/" {
+		t.Fatalf("bad: %v", keys)
+	}
+
+	// Should list a, b, sub/
+	idx, keys, err = store.KVSListKeys("/web/", "/")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if idx != 1002 {
+		t.Fatalf("bad: %v", idx)
+	}
+	if len(keys) != 3 {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[0] != "/web/a" {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[1] != "/web/b" {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[2] != "/web/sub/" {
+		t.Fatalf("bad: %v", keys)
+	}
+
+	// Should list c
+	idx, keys, err = store.KVSListKeys("/web/sub/", "/")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if idx != 1002 {
+		t.Fatalf("bad: %v", idx)
+	}
+	if len(keys) != 1 {
+		t.Fatalf("bad: %v", keys)
+	}
+	if keys[0] != "/web/sub/c" {
+		t.Fatalf("bad: %v", keys)
+	}
+}
+
 func TestKVSDeleteTree(t *testing.T) {
 	store, err := testStateStore()
 	if err != nil {