100 lines
2.1 KiB
Go
100 lines
2.1 KiB
Go
|
package labelutil
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"encoding/csv"
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
|
||
|
|
"github.com/prometheus/common/model"
|
||
|
|
"gopkg.in/yaml.v2"
|
||
|
|
)
|
||
|
|
|
||
|
|
// LabelSet is a labelSet that can be used as a flag.
|
||
|
|
type LabelSet struct {
|
||
|
|
model.LabelSet `yaml:",inline"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// String implements flag.Value
|
||
|
|
// Format: a=1,b=2
|
||
|
|
func (v LabelSet) String() string {
|
||
|
|
if v.LabelSet == nil {
|
||
|
|
return ""
|
||
|
|
}
|
||
|
|
records := make([]string, 0, len(v.LabelSet)>>1)
|
||
|
|
for k, v := range v.LabelSet {
|
||
|
|
records = append(records, string(k)+"="+string(v))
|
||
|
|
}
|
||
|
|
|
||
|
|
var buf bytes.Buffer
|
||
|
|
w := csv.NewWriter(&buf)
|
||
|
|
if err := w.Write(records); err != nil {
|
||
|
|
panic(err)
|
||
|
|
}
|
||
|
|
w.Flush()
|
||
|
|
return "[" + strings.TrimSpace(buf.String()) + "]"
|
||
|
|
}
|
||
|
|
|
||
|
|
// Set implements flag.Value
|
||
|
|
func (v *LabelSet) Set(s string) error {
|
||
|
|
var ss []string
|
||
|
|
n := strings.Count(s, "=")
|
||
|
|
switch n {
|
||
|
|
case 0:
|
||
|
|
return fmt.Errorf("%s must be formatted as key=value", s)
|
||
|
|
case 1:
|
||
|
|
ss = append(ss, strings.Trim(s, `"`))
|
||
|
|
default:
|
||
|
|
r := csv.NewReader(strings.NewReader(s))
|
||
|
|
var err error
|
||
|
|
ss, err = r.Read()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
out := model.LabelSet{}
|
||
|
|
for _, pair := range ss {
|
||
|
|
kv := strings.SplitN(pair, "=", 2)
|
||
|
|
if len(kv) != 2 {
|
||
|
|
return fmt.Errorf("%s must be formatted as key=value", pair)
|
||
|
|
}
|
||
|
|
out[model.LabelName(kv[0])] = model.LabelValue(kv[1])
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := out.Validate(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
v.LabelSet = out
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// UnmarshalYAML the Unmarshaler interface of the yaml pkg.
|
||
|
|
func (v *LabelSet) UnmarshalYAML(unmarshal func(interface{}) error) error {
|
||
|
|
lbSet := model.LabelSet{}
|
||
|
|
err := unmarshal(&lbSet)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
v.LabelSet = lbSet
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// MarshalYAML implements yaml.Marshaller.
|
||
|
|
func (v LabelSet) MarshalYAML() (interface{}, error) {
|
||
|
|
out, err := yaml.Marshal(ModelLabelSetToMap(v.LabelSet))
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return string(out), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// ModelLabelSetToMap convert a model.LabelSet to a map[string]string
|
||
|
|
func ModelLabelSetToMap(m model.LabelSet) map[string]string {
|
||
|
|
result := map[string]string{}
|
||
|
|
for k, v := range m {
|
||
|
|
result[string(k)] = string(v)
|
||
|
|
}
|
||
|
|
return result
|
||
|
|
}
|