traefik/kv.go
2015-10-23 22:21:16 +02:00

196 lines
4.8 KiB
Go

/*
Copyright
*/
package main
import (
"bytes"
"github.com/docker/libkv"
"github.com/docker/libkv/store/boltdb"
"github.com/docker/libkv/store/consul"
"github.com/docker/libkv/store/etcd"
"github.com/docker/libkv/store/zookeeper"
"strings"
"text/template"
"errors"
"github.com/BurntSushi/toml"
"github.com/BurntSushi/ty/fun"
log "github.com/Sirupsen/logrus"
"github.com/docker/libkv/store"
"time"
)
type KvProvider struct {
Watch bool
Endpoint string
Prefix string
Filename string
StoreType store.Backend
kvclient store.Store
}
func NewConsulProvider(provider *ConsulProvider) *KvProvider {
kvProvider := new(KvProvider)
kvProvider.Watch = provider.Watch
kvProvider.Endpoint = provider.Endpoint
kvProvider.Prefix = provider.Prefix
kvProvider.Filename = provider.Filename
kvProvider.StoreType = store.CONSUL
return kvProvider
}
func NewEtcdProvider(provider *EtcdProvider) *KvProvider {
kvProvider := new(KvProvider)
kvProvider.Watch = provider.Watch
kvProvider.Endpoint = provider.Endpoint
kvProvider.Prefix = provider.Prefix
kvProvider.Filename = provider.Filename
kvProvider.StoreType = store.ETCD
return kvProvider
}
func NewZkProvider(provider *ZookepperProvider) *KvProvider {
kvProvider := new(KvProvider)
kvProvider.Watch = provider.Watch
kvProvider.Endpoint = provider.Endpoint
kvProvider.Prefix = provider.Prefix
kvProvider.Filename = provider.Filename
kvProvider.StoreType = store.ZK
return kvProvider
}
func NewBoltDbProvider(provider *BoltDbProvider) *KvProvider {
kvProvider := new(KvProvider)
kvProvider.Watch = provider.Watch
kvProvider.Endpoint = provider.Endpoint
kvProvider.Prefix = provider.Prefix
kvProvider.Filename = provider.Filename
kvProvider.StoreType = store.BOLTDB
return kvProvider
}
func (provider *KvProvider) provide(configurationChan chan<- configMessage) error {
switch provider.StoreType {
case store.CONSUL:
consul.Register()
case store.ETCD:
etcd.Register()
case store.ZK:
zookeeper.Register()
case store.BOLTDB:
boltdb.Register()
default:
return errors.New("Invalid kv store: " + string(provider.StoreType))
}
kv, err := libkv.NewStore(
provider.StoreType,
[]string{provider.Endpoint},
&store.Config{
ConnectionTimeout: 30 * time.Second,
Bucket: "traefik",
},
)
if err != nil {
return err
}
if _, err := kv.List(""); err != nil {
return err
}
provider.kvclient = kv
if provider.Watch {
stopCh := make(chan struct{})
chanKeys, err := kv.WatchTree(provider.Prefix, stopCh)
if err != nil {
return err
}
go func() {
for {
<-chanKeys
configuration := provider.loadConfig()
if configuration != nil {
configurationChan <- configMessage{string(provider.StoreType), configuration}
}
defer close(stopCh)
}
}()
}
configuration := provider.loadConfig()
configurationChan <- configMessage{string(provider.StoreType), configuration}
return nil
}
func (provider *KvProvider) loadConfig() *Configuration {
configuration := new(Configuration)
templateObjects := struct {
Prefix string
}{
provider.Prefix,
}
var KvFuncMap = template.FuncMap{
"List": func(keys ...string) []string {
joinedKeys := strings.Join(keys, "")
keysPairs, err := provider.kvclient.List(joinedKeys)
if err != nil {
log.Error("Error getting keys: ", joinedKeys, err)
return nil
}
directoryKeys := make(map[string]string)
for _, key := range keysPairs {
directory := strings.Split(strings.TrimPrefix(key.Key, strings.TrimPrefix(joinedKeys, "/")), "/")[0]
directoryKeys[directory] = joinedKeys + directory
}
return fun.Values(directoryKeys).([]string)
},
"Get": func(keys ...string) string {
joinedKeys := strings.Join(keys, "")
keyPair, err := provider.kvclient.Get(joinedKeys)
if err != nil {
log.Debug("Error getting key: ", joinedKeys, err)
return ""
} else if keyPair == nil {
return ""
}
return string(keyPair.Value)
},
"Last": func(key string) string {
splittedKey := strings.Split(key, "/")
return splittedKey[len(splittedKey)-1]
},
}
tmpl := template.New(provider.Filename).Funcs(KvFuncMap)
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("templates/kv.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 kv template:", err)
return nil
}
if _, err := toml.Decode(buffer.String(), configuration); err != nil {
log.Error("Error creating kv configuration:", err)
log.Error(buffer.String())
return nil
}
return configuration
}