github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/flagext/labelset.go (about)

     1  package flagext
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/csv"
     6  	"fmt"
     7  	"strings"
     8  
     9  	"github.com/prometheus/common/model"
    10  	"gopkg.in/yaml.v2"
    11  
    12  	"github.com/grafana/loki/pkg/util"
    13  )
    14  
    15  // LabelSet is a labelSet that can be used as a flag.
    16  type LabelSet struct {
    17  	model.LabelSet `yaml:",inline"`
    18  }
    19  
    20  // String implements flag.Value
    21  // Format: a=1,b=2
    22  func (v LabelSet) String() string {
    23  	if v.LabelSet == nil {
    24  		return ""
    25  	}
    26  	records := make([]string, 0, len(v.LabelSet)>>1)
    27  	for k, v := range v.LabelSet {
    28  		records = append(records, string(k)+"="+string(v))
    29  	}
    30  
    31  	var buf bytes.Buffer
    32  	w := csv.NewWriter(&buf)
    33  	if err := w.Write(records); err != nil {
    34  		panic(err)
    35  	}
    36  	w.Flush()
    37  	return "[" + strings.TrimSpace(buf.String()) + "]"
    38  }
    39  
    40  // Set implements flag.Value
    41  func (v *LabelSet) Set(s string) error {
    42  	var ss []string
    43  	n := strings.Count(s, "=")
    44  	switch n {
    45  	case 0:
    46  		return fmt.Errorf("%s must be formatted as key=value", s)
    47  	case 1:
    48  		ss = append(ss, strings.Trim(s, `"`))
    49  	default:
    50  		r := csv.NewReader(strings.NewReader(s))
    51  		var err error
    52  		ss, err = r.Read()
    53  		if err != nil {
    54  			return err
    55  		}
    56  	}
    57  
    58  	out := model.LabelSet{}
    59  	for _, pair := range ss {
    60  		kv := strings.SplitN(pair, "=", 2)
    61  		if len(kv) != 2 {
    62  			return fmt.Errorf("%s must be formatted as key=value", pair)
    63  		}
    64  		out[model.LabelName(kv[0])] = model.LabelValue(kv[1])
    65  	}
    66  
    67  	if err := out.Validate(); err != nil {
    68  		return err
    69  	}
    70  	v.LabelSet = out
    71  	return nil
    72  }
    73  
    74  // UnmarshalYAML the Unmarshaler interface of the yaml pkg.
    75  func (v *LabelSet) UnmarshalYAML(unmarshal func(interface{}) error) error {
    76  	lbSet := model.LabelSet{}
    77  	err := unmarshal(&lbSet)
    78  	if err != nil {
    79  		return err
    80  	}
    81  	v.LabelSet = lbSet
    82  	return nil
    83  }
    84  
    85  // MarshalYAML implements yaml.Marshaller.
    86  func (v LabelSet) MarshalYAML() (interface{}, error) {
    87  	out, err := yaml.Marshal(util.ModelLabelSetToMap(v.LabelSet))
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  	return string(out), nil
    92  }