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

     1  package cfg
     2  
     3  import (
     4  	"flag"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"github.com/awesome-flow/flow/pkg/types"
     9  )
    10  
    11  var (
    12  	cfgFile    string
    13  	pluginPath string
    14  )
    15  
    16  // Redefined in tests
    17  var regFlags = func(cp *CliProvider) {
    18  	if !flag.Parsed() {
    19  		flag.StringVar(&cfgFile, CfgPathKey, "", "Config file path")
    20  		flag.StringVar(&pluginPath, PluginPathKey, "", "Plugin folder path")
    21  		flag.Var(cp, "o", "Extra options")
    22  		flag.Parse()
    23  		if len(cfgFile) > 0 {
    24  			cp.registry[CfgPathKey] = cfgFile
    25  		}
    26  		if len(pluginPath) > 0 {
    27  			cp.registry[PluginPathKey] = pluginPath
    28  		}
    29  	}
    30  }
    31  
    32  // CliProvider serves command-line flag values. By default, it registers a few
    33  // basic flags, backing a full range of config keys by -o attribute.
    34  type CliProvider struct {
    35  	weight   int
    36  	registry map[string]types.Value
    37  	ready    chan struct{}
    38  }
    39  
    40  var _ Provider = (*CliProvider)(nil)
    41  var _ flag.Value = (*CliProvider)(nil)
    42  
    43  // NewCliProvider returns a new instance of CliProvider.
    44  func NewCliProvider(repo *Repository, weight int) (*CliProvider, error) {
    45  	prov := &CliProvider{
    46  		weight:   weight,
    47  		registry: make(map[string]types.Value),
    48  		ready:    make(chan struct{}),
    49  	}
    50  	repo.RegisterProvider(prov)
    51  
    52  	return prov, nil
    53  }
    54  
    55  // Name returns provider name: cli
    56  func (cp *CliProvider) Name() string { return "cli" }
    57  
    58  // Depends returns the list of provider dependencies: default
    59  func (cp *CliProvider) Depends() []string { return []string{"default"} }
    60  
    61  //Weight returns the provider weight
    62  func (cp *CliProvider) Weight() int { return cp.weight }
    63  
    64  // String satisfies flag.Value() interface
    65  func (cp *CliProvider) String() string { return fmt.Sprintf("%v", cp.registry) }
    66  
    67  // Set satisfies flag.Value() interface
    68  func (cp *CliProvider) Set(val string) error {
    69  	if chunks := strings.Split(val, "="); len(chunks) > 2 {
    70  		return fmt.Errorf("Possibly malformed flag (way too many `=`): %q", val)
    71  	} else if len(chunks) == 2 {
    72  		cp.registry[chunks[0]] = chunks[1]
    73  	} else {
    74  		cp.registry[val] = true
    75  	}
    76  	return nil
    77  }
    78  
    79  // SetUp registers a bunch of command line flags (if not registered).
    80  // Flag list:
    81  // * -config.path: the config file location
    82  // * -plugins.path: the plugin folder location
    83  // * -o: extra options, ex: -o system.maxproc=4 -o pipeline.tcp_rcv.connect=udp
    84  func (cp *CliProvider) SetUp(repo *Repository) error {
    85  	defer close(cp.ready)
    86  	regFlags(cp)
    87  	for k := range cp.registry {
    88  		if err := repo.RegisterKey(types.NewKey(k), cp); err != nil {
    89  			return err
    90  		}
    91  	}
    92  	return nil
    93  }
    94  
    95  // TearDown is a no-op operation for CliProvider
    96  func (cp *CliProvider) TearDown(*Repository) error { return nil }
    97  
    98  // Get is the primary method for fetching values from the cli registry
    99  func (cp *CliProvider) Get(key types.Key) (*types.KeyValue, bool) {
   100  	<-cp.ready
   101  	if v, ok := cp.registry[key.String()]; ok {
   102  		return &types.KeyValue{Key: key, Value: v}, ok
   103  	}
   104  	return nil, false
   105  }