github.com/Be-MobileNV/viper@v1.0.2-0.20190415120027-e0a6cbb1bd49/remote/remote.go (about)

     1  // Copyright © 2015 Steve Francia <spf@spf13.com>.
     2  //
     3  // Use of this source code is governed by an MIT-style
     4  // license that can be found in the LICENSE file.
     5  
     6  // Package remote integrates the remote features of Viper.
     7  package remote
     8  
     9  import (
    10  	"bytes"
    11  	"context"
    12  	"fmt"
    13  	"io"
    14  	"strings"
    15  	"time"
    16  
    17  	"github.com/Be-MobileNV/viper"
    18  	"go.etcd.io/etcd/clientv3"
    19  )
    20  
    21  type remoteConfigProvider struct {
    22  	client *clientv3.Client
    23  }
    24  
    25  func (rc remoteConfigProvider) Get(rp viper.RemoteProvider) (io.Reader, error) {
    26  	if rc.client == nil {
    27  		var err error
    28  		rc.client, err = getConfigManager(rp)
    29  		if err != nil {
    30  			return nil, err
    31  		}
    32  	}
    33  	r, err := rc.client.Get(context.Background(), rp.Path())
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	if len(r.Kvs) < 1 {
    39  		return nil, fmt.Errorf("no key-value pairs matched the request")
    40  	}
    41  
    42  	return bytes.NewReader(r.Kvs[0].Value), nil
    43  }
    44  
    45  func (rc remoteConfigProvider) Watch(rp viper.RemoteProvider) (io.Reader, error) {
    46  	if rc.client == nil {
    47  		var err error
    48  		rc.client, err = getConfigManager(rp)
    49  		if err != nil {
    50  			return nil, err
    51  		}
    52  	}
    53  	w := rc.client.Watch(context.Background(), rp.Path())
    54  	resp := <-w
    55  	if resp.Err() != nil {
    56  		return nil, resp.Err()
    57  	}
    58  	if len(resp.Events) < 1 {
    59  		return nil, fmt.Errorf("no watch events found")
    60  	}
    61  	val := resp.Events[0].Kv.Value
    62  	return bytes.NewReader(val), nil
    63  }
    64  
    65  func (rc remoteConfigProvider) WatchChannel(rp viper.RemoteProvider) (responseChannel <-chan *viper.RemoteResponse, quitwc chan bool, err error) {
    66  	quitwc = make(chan bool)
    67  	viperResponsCh := make(chan *viper.RemoteResponse)
    68  	if rc.client == nil {
    69  		var err error
    70  		rc.client, err = getConfigManager(rp)
    71  		if err != nil {
    72  			return nil, nil, err
    73  		}
    74  	}
    75  	ctx, cancel := context.WithCancel(context.Background())
    76  	w := rc.client.Watch(ctx, rp.Path())
    77  	go func(etcdResponseChannel clientv3.WatchChan, vr chan<- *viper.RemoteResponse, quitwc <-chan bool, cancel context.CancelFunc) {
    78  		for {
    79  			select {
    80  			case <-quitwc:
    81  				cancel()
    82  				return
    83  			case resp := <-etcdResponseChannel:
    84  
    85  				if err := resp.Err(); err != nil {
    86  					vr <- &viper.RemoteResponse{
    87  						Error: err,
    88  					}
    89  					continue
    90  				}
    91  
    92  				if len(resp.Events) < 1 {
    93  					vr <- &viper.RemoteResponse{
    94  						Error: fmt.Errorf("no watch events found"),
    95  					}
    96  					continue
    97  				}
    98  
    99  				vr <- &viper.RemoteResponse{
   100  					Value: resp.Events[0].Kv.Value,
   101  				}
   102  
   103  			}
   104  
   105  		}
   106  	}(w, viperResponsCh, quitwc, cancel)
   107  	return viperResponsCh, quitwc, nil
   108  }
   109  
   110  func getConfigManager(rp viper.RemoteProvider) (*clientv3.Client, error) {
   111  	etcdConfig := clientv3.Config{
   112  		Endpoints:   strings.Split(rp.Endpoint(), ","),
   113  		DialTimeout: 5 * time.Second,
   114  		Username:    rp.Username(),
   115  		Password:    rp.Password(),
   116  	}
   117  	client, err := clientv3.New(etcdConfig)
   118  	if err != nil {
   119  		return nil, err
   120  	}
   121  	return client, nil
   122  }
   123  
   124  func init() {
   125  	viper.RemoteConfig = &remoteConfigProvider{}
   126  }