traefik/consul.go
2015-09-21 18:05:56 +02:00

155 lines
3.6 KiB
Go

package main
import (
"github.com/hashicorp/consul/api"
"text/template"
"bytes"
"github.com/BurntSushi/toml"
"strings"
"github.com/BurntSushi/ty/fun"
"net/http"
)
type Key struct {
Value string
}
type ConsulProvider struct {
Watch bool
Endpoint string
Prefix string
Filename string
consulClient *api.Client
}
var kvClient *api.KV
var ConsulFuncMap = template.FuncMap{
"List": func(keys ...string) []string {
joinedKeys := strings.Join(keys, "")
keysPairs, _, err := kvClient.Keys(joinedKeys, "/", nil)
if err != nil {
log.Error("Error getting keys ", joinedKeys, err)
return nil
}
keysPairs = fun.Filter(func(key string) bool {
if (key == joinedKeys) {
return false
}
return true
}, keysPairs).([]string)
return keysPairs
},
"Get": func(keys ...string) string {
joinedKeys := strings.Join(keys, "")
keyPair, _, err := kvClient.Get(joinedKeys, nil)
if err != nil {
log.Error("Error getting key ", joinedKeys, err)
return ""
}
return string(keyPair.Value)
},
"Last": func(key string) string {
splittedKey := strings.Split(key, "/")
return splittedKey[len(splittedKey) -2]
},
}
func NewConsulProvider() *ConsulProvider {
consulProvider := new(ConsulProvider)
// default values
consulProvider.Watch = true
consulProvider.Prefix = "traefik"
return consulProvider
}
func (provider *ConsulProvider) Provide(configurationChan chan <- *Configuration) {
config := &api.Config{
Address: provider.Endpoint,
Scheme: "http",
HttpClient: http.DefaultClient,
}
consulClient, _ := api.NewClient(config)
provider.consulClient = consulClient
if provider.Watch {
var waitIndex uint64
keypairs, meta, err := consulClient.KV().Keys("", "", nil)
if keypairs == nil && err == nil {
log.Error("Key was not found.")
}
waitIndex = meta.LastIndex
go func() {
for {
opts := api.QueryOptions{
WaitIndex: waitIndex,
}
keypairs, meta, err := consulClient.KV().Keys("", "", &opts)
if keypairs == nil && err == nil {
log.Error("Key was not found.")
}
waitIndex = meta.LastIndex
configuration := provider.loadConsulConfig()
if configuration != nil {
configurationChan <- configuration
}
}
}()
}
configuration := provider.loadConsulConfig()
configurationChan <- configuration
}
func (provider *ConsulProvider) loadConsulConfig() *Configuration {
configuration := new(Configuration)
services := []*api.CatalogService{}
kvClient = provider.consulClient.KV()
servicesName, _, _ := provider.consulClient.Catalog().Services(nil)
for serviceName, _ := range servicesName {
catalogServices, _, _ := provider.consulClient.Catalog().Service(serviceName, "", nil)
for _, catalogService := range catalogServices {
services= append(services, catalogService)
}
}
templateObjects := struct {
Services []*api.CatalogService
}{
services,
}
tmpl := template.New(provider.Filename).Funcs(ConsulFuncMap)
if len(provider.Filename) > 0 {
_, err := tmpl.ParseFiles(provider.Filename)
if err != nil {
log.Error("Error reading file", err)
return nil
}
} else {
buf, err := Asset("providerTemplates/consul.tmpl")
if err != nil {
log.Error("Error reading file", err)
}
_, err = tmpl.Parse(string(buf))
if err != nil {
log.Error("Error reading file", err)
return nil
}
}
var buffer bytes.Buffer
err := tmpl.Execute(&buffer, templateObjects)
if err != nil {
log.Error("Error with consul template:", err)
return nil
}
if _, err := toml.Decode(buffer.String(), configuration); err != nil {
log.Error("Error creating consul configuration:", err)
return nil
}
return configuration
}