From 255be3559ccbf8e677e6a34e1e179adefba39e8e Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 12 Dec 2013 10:35:50 -0800
Subject: [PATCH] Support datacenter listing

---
 consul/catalog_endpoint.go      | 16 ++++++++++++++
 consul/catalog_endpoint_test.go | 39 +++++++++++++++++++++++++++++++++
 consul/state_store.go           |  1 +
 3 files changed, 56 insertions(+)

diff --git a/consul/catalog_endpoint.go b/consul/catalog_endpoint.go
index 89cee110fb..065e9ec0bd 100644
--- a/consul/catalog_endpoint.go
+++ b/consul/catalog_endpoint.go
@@ -47,3 +47,19 @@ func (c *Catalog) Deregister(args *rpc.DeregisterRequest, reply *struct{}) error
 	}
 	return nil
 }
+
+// ListDatacenters is used to query for the list of known datacenters
+func (c *Catalog) ListDatacenters(args *struct{}, reply *[]string) error {
+	c.srv.remoteLock.RLock()
+	defer c.srv.remoteLock.RUnlock()
+
+	// Read the known DCs
+	var dcs []string
+	for dc := range c.srv.remoteConsuls {
+		dcs = append(dcs, dc)
+	}
+
+	// Return
+	*reply = dcs
+	return nil
+}
diff --git a/consul/catalog_endpoint_test.go b/consul/catalog_endpoint_test.go
index 4a9c47ebb3..3be0645523 100644
--- a/consul/catalog_endpoint_test.go
+++ b/consul/catalog_endpoint_test.go
@@ -5,6 +5,7 @@ import (
 	"github.com/hashicorp/consul/rpc"
 	nrpc "net/rpc"
 	"os"
+	"sort"
 	"testing"
 	"time"
 )
@@ -144,3 +145,41 @@ func TestCatalogDeregister(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 }
+
+func TestCatalogListDatacenters(t *testing.T) {
+	dir1, s1 := testServer(t)
+	defer os.RemoveAll(dir1)
+	defer s1.Shutdown()
+	client := rpcClient(t, s1)
+	defer client.Close()
+
+	dir2, s2 := testServerDC(t, "dc2")
+	defer os.RemoveAll(dir2)
+	defer s2.Shutdown()
+
+	// Try to join
+	addr := fmt.Sprintf("127.0.0.1:%d",
+		s1.config.SerfWANConfig.MemberlistConfig.Port)
+	if err := s2.JoinWAN(addr); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	time.Sleep(10 * time.Millisecond)
+
+	var out []string
+	if err := client.Call("Catalog.ListDatacenters", struct{}{}, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Sort the dcs
+	sort.Strings(out)
+
+	if len(out) != 2 {
+		t.Fatalf("bad: %v", out)
+	}
+	if out[0] != "dc1" {
+		t.Fatalf("bad: %v", out)
+	}
+	if out[1] != "dc2" {
+		t.Fatalf("bad: %v", out)
+	}
+}
diff --git a/consul/state_store.go b/consul/state_store.go
index cd18c82f5b..761fd476e8 100644
--- a/consul/state_store.go
+++ b/consul/state_store.go
@@ -82,6 +82,7 @@ func (s *StateStore) initialize() error {
 		`CREATE TABLE nodes (name text unique, address text);`,
 		`CREATE TABLE services (node text REFERENCES nodes(name) ON DELETE CASCADE, service text, tag text, port integer);`,
 		`CREATE INDEX servName ON services(service);`,
+		`CREATE INDEX nodeName ON services(node);`,
 	}
 	for _, t := range tables {
 		if _, err := s.db.Exec(t); err != nil {