From 3a203185d75204dfc8c797f22fecbc13085498b4 Mon Sep 17 00:00:00 2001
From: Kelsey Hightower <kelsey.hightower@gmail.com>
Date: Mon, 26 May 2014 08:24:22 -0700
Subject: [PATCH] Add support for consul iteration

---
 .gitignore                                    |  1 -
 confd.go                                      | 13 ++------
 config/config.go                              |  9 ++++++
 config/config_test.go                         |  1 +
 env/client.go                                 | 12 +++----
 etcd/etcdutil/client.go                       |  3 +-
 etcd/etcdutil/client_test.go                  |  3 +-
 integration/confdir/conf.d/basic.toml         |  9 ++++++
 integration/confdir/conf.d/iteration.toml     |  6 ++++
 integration/confdir/templates/basic.conf.tmpl |  5 +++
 .../confdir/templates/iteration.conf.tmpl     | 17 ++++++++++
 integration/consul/test.sh                    | 18 +++++++++++
 integration/etcd/test.sh                      | 15 +++++++++
 node/node.go                                  | 31 +++++++++++++++++++
 resource/template/resource.go                 | 25 +++++++++++++++
 15 files changed, 147 insertions(+), 21 deletions(-)
 create mode 100644 integration/confdir/conf.d/basic.toml
 create mode 100644 integration/confdir/conf.d/iteration.toml
 create mode 100644 integration/confdir/templates/basic.conf.tmpl
 create mode 100644 integration/confdir/templates/iteration.conf.tmpl
 create mode 100644 integration/consul/test.sh
 create mode 100644 integration/etcd/test.sh
 create mode 100644 node/node.go

diff --git a/.gitignore b/.gitignore
index 7ba751459..e69de29bb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +0,0 @@
-confd
diff --git a/confd.go b/confd.go
index 0c042245c..759b03d75 100644
--- a/confd.go
+++ b/confd.go
@@ -4,11 +4,11 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"os"
 	"strings"
 	"time"
-	"errors"
 
 	"github.com/kelseyhightower/confd/config"
 	"github.com/kelseyhightower/confd/consul"
@@ -22,13 +22,11 @@ var (
 	configFile        = ""
 	defaultConfigFile = "/etc/confd/confd.toml"
 	onetime           bool
-	backend           = ""
 )
 
 func init() {
 	flag.StringVar(&configFile, "config-file", "", "the confd config file")
 	flag.BoolVar(&onetime, "onetime", false, "run once and exit")
-	flag.StringVar(&backend, "backend", "", "backend to use")
 }
 
 func main() {
@@ -54,7 +52,7 @@ func main() {
 	log.Notice("Starting confd")
 
 	// Create the storage client
-	store, err := createStoreClient(backend)
+	store, err := createStoreClient(config.Backend())
 	if err != nil {
 		log.Fatal(err.Error())
 	}
@@ -76,13 +74,6 @@ func main() {
 // createStoreClient is used to create a storage client based
 // on our configuration. Either an etcd or Consul client.
 func createStoreClient(backend string) (template.StoreClient, error) {
-	if backend == "" {
-		if config.Consul() {
-			backend = "consul"
-		} else {
-			backend = "etcd"
-		}
-	}
 	log.Notice("Backend set to " + backend)
 	switch backend {
 	case "consul":
diff --git a/config/config.go b/config/config.go
index 41ccd5982..fc0d21ff3 100644
--- a/config/config.go
+++ b/config/config.go
@@ -17,6 +17,7 @@ import (
 )
 
 var (
+	backend      string
 	clientCert   string
 	clientKey    string
 	clientCaKeys string
@@ -42,6 +43,7 @@ type Config struct {
 
 // confd represents the parsed configuration settings.
 type confd struct {
+	Backend      string   `toml:"backend"`
 	Debug        bool     `toml:"debug"`
 	ClientCert   string   `toml:"client_cert"`
 	ClientKey    string   `toml:"client_key"`
@@ -60,6 +62,7 @@ type confd struct {
 }
 
 func init() {
+	flag.StringVar(&backend, "backend", "", "backend to use")
 	flag.BoolVar(&debug, "debug", false, "enable debug logging")
 	flag.StringVar(&clientCert, "client-cert", "", "the client cert")
 	flag.StringVar(&clientKey, "client-key", "", "the client key")
@@ -109,6 +112,10 @@ func LoadConfig(path string) error {
 	return nil
 }
 
+func Backend() string {
+	return config.Confd.Backend
+}
+
 // Debug reports whether debug mode is enabled.
 func Debug() bool {
 	return config.Confd.Debug
@@ -276,6 +283,8 @@ func processFlags() {
 
 func setConfigFromFlag(f *flag.Flag) {
 	switch f.Name {
+	case "backend":
+		config.Confd.Backend = backend
 	case "debug":
 		config.Confd.Debug = debug
 	case "client-cert":
diff --git a/config/config_test.go b/config/config_test.go
index 1708ff32a..31d63a253 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -5,6 +5,7 @@ package config
 
 import (
 	"testing"
+
 	"github.com/kelseyhightower/confd/log"
 )
 
diff --git a/env/client.go b/env/client.go
index 931b503d9..5bfa05934 100644
--- a/env/client.go
+++ b/env/client.go
@@ -8,7 +8,7 @@ import (
 var replacer = strings.NewReplacer("/", "_")
 
 // Client provides a wrapper around the consulkv client
-type Client struct {}
+type Client struct{}
 
 // NewEnvClient returns a new client
 func NewEnvClient() (*Client, error) {
@@ -19,18 +19,16 @@ func NewEnvClient() (*Client, error) {
 func (c *Client) GetValues(keys []string) (map[string]interface{}, error) {
 	vars := make(map[string]interface{})
 	for _, key := range keys {
-		k := transform(key)		
+		k := transform(key)
 		value := os.Getenv(k)
 		if value != "" {
 			vars[key] = value
-		} 
+		}
 	}
 	return vars, nil
 }
 
-func transform(key string) (string) {
+func transform(key string) string {
 	k := strings.TrimPrefix(key, "/")
-	k = replacer.Replace(k)
-	k = strings.ToUpper(k)
-	return k	
+	return strings.ToUpper(replacer.Replace(k))
 }
diff --git a/etcd/etcdutil/client.go b/etcd/etcdutil/client.go
index 2c1801a13..f87596711 100644
--- a/etcd/etcdutil/client.go
+++ b/etcd/etcdutil/client.go
@@ -5,8 +5,9 @@ package etcdutil
 
 import (
 	"errors"
-	"github.com/coreos/go-etcd/etcd"
 	"strings"
+
+	"github.com/coreos/go-etcd/etcd"
 )
 
 // Client is a wrapper around the etcd client
diff --git a/etcd/etcdutil/client_test.go b/etcd/etcdutil/client_test.go
index d5291c651..2aaf06c36 100644
--- a/etcd/etcdutil/client_test.go
+++ b/etcd/etcdutil/client_test.go
@@ -4,9 +4,10 @@
 package etcdutil
 
 import (
+	"testing"
+
 	"github.com/coreos/go-etcd/etcd"
 	"github.com/kelseyhightower/confd/etcd/etcdtest"
-	"testing"
 )
 
 func TestGetValues(t *testing.T) {
diff --git a/integration/confdir/conf.d/basic.toml b/integration/confdir/conf.d/basic.toml
new file mode 100644
index 000000000..5250447f3
--- /dev/null
+++ b/integration/confdir/conf.d/basic.toml
@@ -0,0 +1,9 @@
+[template]
+src = "basic.conf.tmpl"
+dest = "/tmp/confd-basic-test.conf"
+keys = [
+  "/database/host",
+  "/database/password",
+  "/database/port",
+  "/database/username",
+]
diff --git a/integration/confdir/conf.d/iteration.toml b/integration/confdir/conf.d/iteration.toml
new file mode 100644
index 000000000..94b027d5b
--- /dev/null
+++ b/integration/confdir/conf.d/iteration.toml
@@ -0,0 +1,6 @@
+[template]
+src = "iteration.conf.tmpl"
+dest = "/tmp/confd-iteration-test.conf"
+keys = [
+  "/upstream",
+]
diff --git a/integration/confdir/templates/basic.conf.tmpl b/integration/confdir/templates/basic.conf.tmpl
new file mode 100644
index 000000000..d8652e774
--- /dev/null
+++ b/integration/confdir/templates/basic.conf.tmpl
@@ -0,0 +1,5 @@
+[database]
+host={{.database_host}}
+password={{.database_password}}
+port={{.database_port}}
+username={{.database_username}}
diff --git a/integration/confdir/templates/iteration.conf.tmpl b/integration/confdir/templates/iteration.conf.tmpl
new file mode 100644
index 000000000..0dc59bd89
--- /dev/null
+++ b/integration/confdir/templates/iteration.conf.tmpl
@@ -0,0 +1,17 @@
+upstream app {
+{{range $server := GetDir "upstream"}}
+    server {{$server.Value}};
+{{end}}
+}
+
+server {
+    server_name  www.example.com;
+
+    location / {
+        proxy_pass        http://app;
+        proxy_redirect    off;
+        proxy_set_header  Host             $host;
+        proxy_set_header  X-Real-IP        $remote_addr;
+        proxy_set_header  X-Forwarded-For  $proxy_add_x_forwarded_for;
+   }
+}
diff --git a/integration/consul/test.sh b/integration/consul/test.sh
new file mode 100644
index 000000000..4b5e8b5a8
--- /dev/null
+++ b/integration/consul/test.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+docker_host=${1}
+consul_id=$(docker run -p 8500:8500 -p 8400:8400 -d kelseyhightower/consul -server -bootstrap -ui-dir /opt/consul/web_ui -client 0.0.0.0 -advertise ${docker_host})
+
+# Configure consul
+curl -X PUT http://${docker_host}:8500/v1/kv/database/host -d '127.0.0.1'
+curl -X PUT http://${docker_host}:8500/v1/kv/database/password -d 'p@sSw0rd'
+curl -X PUT http://${docker_host}:8500/v1/kv/database/port -d '3306'
+curl -X PUT http://${docker_host}:8500/v1/kv/database/username -d 'confd'
+curl -X PUT http://${docker_host}:8500/v1/kv/upstream/app1 -d '10.0.1.10:8080'
+curl -X PUT http://${docker_host}:8500/v1/kv/upstream/app2 -d '10.0.1.11:8080'
+
+# Run confd
+./confd -onetime -confdir ./integration/confdir -backend consul -consul-addr "${1}:8500"
+
+# Cleanup
+docker stop $consul_id
diff --git a/integration/etcd/test.sh b/integration/etcd/test.sh
new file mode 100644
index 000000000..0c13658a7
--- /dev/null
+++ b/integration/etcd/test.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+etcd_host=${1}
+etcd_port=${2}
+
+# Configure consul
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/database/host -d value=127.0.0.1
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/database/password -d value=p@sSw0rd
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/database/port -d value=3306
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/database/username -d value=confd
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/upstream/app1 -d value=10.0.1.10:8080
+curl -L -X PUT http://${etcd_host}:${etcd_port}/v2/keys/upstream/app2 -d value=10.0.1.11:8080
+
+# Run confd
+./confd -onetime -confdir ./integration/confdir -backend etcd -node "http://${etcd_host}:${etcd_port}"
diff --git a/node/node.go b/node/node.go
new file mode 100644
index 000000000..07099f71f
--- /dev/null
+++ b/node/node.go
@@ -0,0 +1,31 @@
+package node
+
+type Node struct {
+	Key   string
+	Value interface{}
+}
+
+type Directory map[string][]Node
+
+func NewDirectory() Directory {
+	d := make(Directory)
+	return d
+}
+
+// Add.
+func (d Directory) Add(key string, node Node) {
+	if d[key] == nil {
+		d[key] = make([]Node, 0)
+	}
+	d[key] = append(d[key], node)
+}
+
+// Get.
+func (d Directory) Get(key string) []Node {
+	return d[key]
+}
+
+// Set.
+func (d Directory) Set(key string, nodes []Node) {
+	d[key] = nodes
+}
diff --git a/resource/template/resource.go b/resource/template/resource.go
index dcc8f08ac..0e3277553 100644
--- a/resource/template/resource.go
+++ b/resource/template/resource.go
@@ -19,6 +19,7 @@ import (
 	"github.com/BurntSushi/toml"
 	"github.com/kelseyhightower/confd/config"
 	"github.com/kelseyhightower/confd/log"
+	"github.com/kelseyhightower/confd/node"
 )
 
 var replacer = strings.NewReplacer("/", "_")
@@ -41,6 +42,7 @@ type TemplateResource struct {
 	StageFile   *os.File
 	Src         string
 	Vars        map[string]interface{}
+	Dirs        node.Directory
 	storeClient StoreClient
 }
 
@@ -70,10 +72,31 @@ func (t *TemplateResource) setVars() error {
 	if err != nil {
 		return err
 	}
+	t.setDirs(vars)
 	t.Vars = cleanKeys(vars, config.Prefix())
 	return nil
 }
 
+// setDirs sets the Dirs for the template resource.
+// All keys are grouped based on their directory path names.
+// For example, /upstream/app1 and upstream/app2 will be grouped as
+//    {
+//        "upstream": []Node{
+//            {"app1": value}},
+//            {"app2": value}},
+//         }
+//    }
+//
+// Dirs are exposed to resource templated to enable iteration.
+func (t *TemplateResource) setDirs(vars map[string]interface{}) {
+	d := node.NewDirectory()
+	for k, v := range vars {
+		directory := filepath.Dir(filepath.Join("/", strings.TrimPrefix(k, config.Prefix())))
+		d.Add(pathToKey(directory, config.Prefix()), node.Node{filepath.Base(k), v})
+	}
+	t.Dirs = d
+}
+
 // createStageFile stages the src configuration file by processing the src
 // template and setting the desired owner, group, and mode. It also sets the
 // StageFile for the template resource.
@@ -93,6 +116,8 @@ func (t *TemplateResource) createStageFile() error {
 	log.Debug("Compiling source template " + t.Src)
 	tplFuncMap := make(template.FuncMap)
 	tplFuncMap["Base"] = path.Base
+
+	tplFuncMap["GetDir"] = t.Dirs.Get
 	tmpl := template.Must(template.New(path.Base(t.Src)).Funcs(tplFuncMap).ParseFiles(t.Src))
 	if err = tmpl.Execute(temp, t.Vars); err != nil {
 		return err