Merge pull request #22 from heetch/batch

Prefetch options for etcd and Consul

.travis.yml

@@ -12,7 +12,7 @@ before_install:
   - export PATH="$GOPATH/bin:$PATH"
   - docker run -d -p 2379:2379 quay.io/coreos/etcd /usr/local/bin/etcd -advertise-client-urls http://0.0.0.0:2379 -listen-client-urls http://0.0.0.0:2379
   - docker run -d -p 8500:8500 --name consul consul
-  - docker run -d -p 8200:8200 --cap-add=IPC_LOCK -e 'VAULT_DEV_ROOT_TOKEN_ID=root' -e 'VAULT_DEV_LISTEN_ADDRESS=0.0.0.0:8200' vault
+  - docker run -d -p 8200:8200 --cap-add=IPC_LOCK -e 'VAULT_DEV_ROOT_TOKEN_ID=root' -e 'VAULT_DEV_LISTEN_ADDRESS=0.0.0.0:8200' vault:0.9.6
 
 env:
   - VAULT_ADDR=http://127.0.0.1:8200

Gopkg.toml

@@ -20,7 +20,7 @@
 
 [[constraint]]
   name = "github.com/hashicorp/vault"
-  version = "0.9.3"
+  version = "0.9.0"
 
 [prune]
   go-tests = true

backend/consul/consul.go

@@ -3,27 +3,46 @@ package consul
 import (
 	"context"
 	"path"
+	"strings"
 
 	"github.com/hashicorp/consul/api"
 	"github.com/heetch/confita/backend"
 )
 
 // Backend loads keys from Consul.
 type Backend struct {
-	client *api.Client
-	prefix string
+	client   *api.Client
+	prefix   string
+	prefetch bool
+	cache    map[string][]byte
 }
 
 // NewBackend creates a configuration loader that loads from Consul.
-func NewBackend(client *api.Client, prefix string) *Backend {
-	return &Backend{
+func NewBackend(client *api.Client, opts ...Option) *Backend {
+	b := Backend{
 		client: client,
-		prefix: prefix,
 	}
+
+	for _, opt := range opts {
+		opt(&b)
+	}
+
+	return &b
 }
 
 // Get loads the given key from Consul.
 func (b *Backend) Get(ctx context.Context, key string) ([]byte, error) {
+	if b.cache == nil && b.prefetch {
+		err := b.fetchTree(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if b.cache != nil {
+		return b.fromCache(ctx, key)
+	}
+
 	var opt api.QueryOptions
 
 	kv, _, err := b.client.KV().Get(path.Join(b.prefix, key), opt.WithContext(ctx))
@@ -38,7 +57,52 @@ func (b *Backend) Get(ctx context.Context, key string) ([]byte, error) {
 	return kv.Value, nil
 }
 
+func (b *Backend) fetchTree(ctx context.Context) error {
+	var opt api.QueryOptions
+
+	list, _, err := b.client.KV().List(b.prefix, opt.WithContext(ctx))
+	if err != nil {
+		return err
+	}
+
+	b.cache = make(map[string][]byte)
+
+	for _, kv := range list {
+		b.cache[strings.TrimLeft(kv.Key, b.prefix+"/")] = kv.Value
+	}
+
+	return nil
+}
+
+func (b *Backend) fromCache(ctx context.Context, key string) ([]byte, error) {
+	v, ok := b.cache[key]
+	if !ok {
+		return nil, backend.ErrNotFound
+	}
+
+	return v, nil
+}
+
 // Name returns the name of the backend.
 func (b *Backend) Name() string {
 	return "consul"
 }
+
+// Option is used to configure the Consul backend.
+type Option func(*Backend)
+
+// WithPrefix is used to specify the prefix to prepend on every keys.
+func WithPrefix(prefix string) Option {
+	return func(b *Backend) {
+		b.prefix = prefix
+	}
+}
+
+// WithPrefetch is used to prefetch the entire tree and cache it to
+// avoid further round trips. If the WithPrefix option is used, will fetch
+// the tree under the specified prefix.
+func WithPrefetch() Option {
+	return func(b *Backend) {
+		b.prefetch = true
+	}
+}

backend/consul/consul_test.go

@@ -16,7 +16,7 @@ func TestConsulBackend(t *testing.T) {
 	require.NoError(t, err)
 	defer client.KV().DeleteTree(prefix, nil)
 
-	b := NewBackend(client, prefix)
+	b := NewBackend(client, WithPrefix(prefix))
 
 	t.Run("NotFound", func(t *testing.T) {
 		_, err = b.Get(context.Background(), "something that doesn't exist")
@@ -42,3 +42,45 @@ func TestConsulBackend(t *testing.T) {
 		require.Error(t, err)
 	})
 }
+
+func TestConsulBackendWithPrefetch(t *testing.T) {
+	prefix := "confita-tests"
+
+	client, err := api.NewClient(api.DefaultConfig())
+	require.NoError(t, err)
+	defer client.KV().DeleteTree(prefix, nil)
+
+	b := NewBackend(client, WithPrefix(prefix), WithPrefetch())
+
+	t.Run("OK", func(t *testing.T) {
+		_, err = client.KV().Put(&api.KVPair{Key: prefix + "/key1", Value: []byte("value1")}, nil)
+		require.NoError(t, err)
+
+		_, err = client.KV().Put(&api.KVPair{Key: prefix + "/key2", Value: []byte("value2")}, nil)
+		require.NoError(t, err)
+
+		_, err = client.KV().Put(&api.KVPair{Key: prefix + "/key3", Value: []byte("value3")}, nil)
+		require.NoError(t, err)
+
+		val, err := b.Get(context.Background(), "key1")
+		require.NoError(t, err)
+
+		// deleting the tree
+		client.KV().DeleteTree(prefix, nil)
+
+		// WithPrefetch should have prefetched all the keys
+		// they should be available even if the tree has been removed.
+		val, err = b.Get(context.Background(), "key1")
+		require.NoError(t, err)
+		require.Equal(t, []byte("value1"), val)
+
+		val, err = b.Get(context.Background(), "key2")
+		require.NoError(t, err)
+		require.Equal(t, []byte("value2"), val)
+
+		val, err = b.Get(context.Background(), "key3")
+		require.NoError(t, err)
+		require.Equal(t, []byte("value3"), val)
+	})
+
+}

backend/etcd/etcd.go

@@ -3,27 +3,46 @@ package etcd
 import (
 	"context"
 	"path"
+	"strings"
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/heetch/confita/backend"
 )
 
 // Backend loads keys from etcd.
 type Backend struct {
-	client *clientv3.Client
-	prefix string
+	client   *clientv3.Client
+	prefix   string
+	prefetch bool
+	cache    map[string][]byte
 }
 
 // NewBackend creates a configuration loader that loads from etcd.
-func NewBackend(client *clientv3.Client, prefix string) *Backend {
-	return &Backend{
+func NewBackend(client *clientv3.Client, opts ...Option) *Backend {
+	b := Backend{
 		client: client,
-		prefix: prefix,
 	}
+
+	for _, opt := range opts {
+		opt(&b)
+	}
+
+	return &b
 }
 
 // Get loads the given key from etcd.
 func (b *Backend) Get(ctx context.Context, key string) ([]byte, error) {
+	if b.cache == nil && b.prefetch {
+		err := b.fetchTree(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if b.cache != nil {
+		return b.fromCache(ctx, key)
+	}
+
 	resp, err := b.client.Get(ctx, path.Join(b.prefix, key))
 	if err != nil {
 		return nil, err
@@ -36,7 +55,50 @@ func (b *Backend) Get(ctx context.Context, key string) ([]byte, error) {
 	return resp.Kvs[0].Value, nil
 }
 
+func (b *Backend) fetchTree(ctx context.Context) error {
+	resp, err := b.client.KV.Get(ctx, b.prefix, clientv3.WithPrefix())
+	if err != nil {
+		return err
+	}
+
+	b.cache = make(map[string][]byte)
+
+	for _, kv := range resp.Kvs {
+		b.cache[strings.TrimLeft(string(kv.Key), b.prefix+"/")] = kv.Value
+	}
+
+	return nil
+}
+
+func (b *Backend) fromCache(ctx context.Context, key string) ([]byte, error) {
+	v, ok := b.cache[key]
+	if !ok {
+		return nil, backend.ErrNotFound
+	}
+
+	return v, nil
+}
+
 // Name returns the name of the backend.
 func (b *Backend) Name() string {
 	return "etcd"
 }
+
+// Option is used to configure the Consul backend.
+type Option func(*Backend)
+
+// WithPrefix is used to specify the prefix to prepend on every keys.
+func WithPrefix(prefix string) Option {
+	return func(b *Backend) {
+		b.prefix = prefix
+	}
+}
+
+// WithPrefetch is used to prefetch the entire tree and cache it to
+// avoid further round trips. If the WithPrefix option is used, will fetch
+// the tree under the specified prefix.
+func WithPrefetch() Option {
+	return func(b *Backend) {
+		b.prefetch = true
+	}
+}