github.com/whiteboxio/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/cfg/yaml_provider.go (about)

     1  package cfg
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  
     7  	"github.com/awesome-flow/flow/pkg/types"
     8  	fsnotify "github.com/fsnotify/fsnotify"
     9  	yaml "gopkg.in/yaml.v2"
    10  )
    11  
    12  // Redefined in tests
    13  var readRaw = func(source string) (map[interface{}]interface{}, error) {
    14  	out := make(map[interface{}]interface{})
    15  	data, err := ioutil.ReadFile(source)
    16  	if err != nil {
    17  		return nil, fmt.Errorf("failed to read yaml config file %q: %s", source, err)
    18  	}
    19  	if err := yaml.Unmarshal(data, &out); err != nil {
    20  		return nil, err
    21  	}
    22  	return out, nil
    23  }
    24  
    25  type YamlProvider struct {
    26  	weight   int
    27  	source   string
    28  	options  *YamlProviderOptions
    29  	watcher  *fsnotify.Watcher
    30  	registry map[string]types.Value
    31  	ready    chan struct{}
    32  }
    33  
    34  type YamlProviderOptions struct {
    35  	Watch bool
    36  }
    37  
    38  var _ Provider = (*YamlProvider)(nil)
    39  
    40  func NewYamlProvider(repo *Repository, weight int) (*YamlProvider, error) {
    41  	return NewYamlProviderWithOptions(repo, weight, &YamlProviderOptions{})
    42  }
    43  
    44  func NewYamlProviderWithOptions(repo *Repository, weight int, options *YamlProviderOptions) (*YamlProvider, error) {
    45  	return NewYamlProviderFromSource(repo, weight, options, "")
    46  }
    47  
    48  func NewYamlProviderFromSource(repo *Repository, weight int, options *YamlProviderOptions, source string) (*YamlProvider, error) {
    49  	prov := &YamlProvider{
    50  		source:   source,
    51  		weight:   weight,
    52  		options:  options,
    53  		registry: make(map[string]types.Value),
    54  		ready:    make(chan struct{}),
    55  	}
    56  	repo.RegisterProvider(prov)
    57  	return prov, nil
    58  }
    59  
    60  func (yp *YamlProvider) Name() string      { return "yaml" }
    61  func (yp *YamlProvider) Depends() []string { return []string{"cli", "env"} }
    62  func (yp *YamlProvider) Weight() int       { return yp.weight }
    63  
    64  func (yp *YamlProvider) SetUp(repo *Repository) error {
    65  	defer close(yp.ready)
    66  
    67  	if len(yp.source) == 0 {
    68  		source, ok := repo.Get(types.NewKey(CfgPathKey))
    69  		if !ok {
    70  			return fmt.Errorf("Failed to get yaml config path from repo")
    71  		}
    72  		yp.source = source.(string)
    73  	}
    74  
    75  	// if _, err := os.Stat(yp.source); err != nil {
    76  	// 	return fmt.Errorf("failed to read yaml config %q: %s", yp.source, err)
    77  	// }
    78  
    79  	if yp.options.Watch {
    80  		watcher, err := fsnotify.NewWatcher()
    81  		if err != nil {
    82  			return fmt.Errorf("failed to start a yaml watcher: %s", err)
    83  		}
    84  		if err := watcher.Add(yp.source); err != nil {
    85  			return fmt.Errorf("failed to add a new watchable file %q: %s", yp.source, err)
    86  		}
    87  		yp.watcher = watcher
    88  
    89  		go yp.watch()
    90  	}
    91  
    92  	rawData, err := readRaw(yp.source)
    93  	if err != nil {
    94  		return err
    95  	}
    96  	for k, v := range flatten(rawData) {
    97  		yp.registry[k] = v
    98  		if repo != nil {
    99  			if err := repo.RegisterKey(types.NewKey(k), yp); err != nil {
   100  				return err
   101  			}
   102  		}
   103  	}
   104  
   105  	return nil
   106  }
   107  
   108  func flatten(in map[interface{}]interface{}) map[string]types.Value {
   109  	out := make(map[string]types.Value)
   110  	for k, v := range in {
   111  		if vmap, ok := v.(map[interface{}]interface{}); ok {
   112  			for sk, sv := range flatten(vmap) {
   113  				out[k.(string)+types.KeySepCh+sk] = types.Value(sv)
   114  			}
   115  		} else {
   116  			out[k.(string)] = types.Value(v)
   117  		}
   118  	}
   119  	return out
   120  }
   121  
   122  func (yp *YamlProvider) watch() {
   123  	for event := range yp.watcher.Events {
   124  		fmt.Printf("Received a fsnotify event: %#v", event)
   125  		//TODO (olegs): not implemented
   126  		// if event.Op&fsnotify.Write != 1 {
   127  		// 	continue
   128  		// }
   129  	}
   130  }
   131  
   132  func (yp *YamlProvider) TearDown(repo *Repository) error {
   133  	if yp.watcher != nil {
   134  		if err := yp.watcher.Close(); err != nil {
   135  			return fmt.Errorf("failed to terminate the yaml watcher: %q", err)
   136  		}
   137  	}
   138  	return nil
   139  }
   140  
   141  func (yp *YamlProvider) Get(key types.Key) (*types.KeyValue, bool) {
   142  	<-yp.ready
   143  	if v, ok := yp.registry[key.String()]; ok {
   144  		return &types.KeyValue{Key: key, Value: v}, ok
   145  	}
   146  	return nil, false
   147  }