diff --git a/pkg/promtail/client/client.go b/pkg/promtail/client/client.go index 8296bcedb33e..694491b70138 100644 --- a/pkg/promtail/client/client.go +++ b/pkg/promtail/client/client.go @@ -18,11 +18,10 @@ import ( "github.com/go-kit/kit/log/level" "github.com/gogo/protobuf/proto" "github.com/golang/snappy" - "github.com/prometheus/client_golang/prometheus" - "github.com/prometheus/common/model" - "github.com/grafana/loki/pkg/helpers" "github.com/grafana/loki/pkg/logproto" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/common/model" ) const contentType = "application/x-protobuf" @@ -83,7 +82,7 @@ func New(cfg Config, logger log.Logger) Client { quit: make(chan struct{}), entries: make(chan entry), - externalLabels: cfg.ExternalLabels, + externalLabels: cfg.ExternalLabels.LabelSet, } c.wg.Add(1) go c.run() diff --git a/pkg/promtail/client/config.go b/pkg/promtail/client/config.go index 5bd38d2f9bec..1cb129d7b4c6 100644 --- a/pkg/promtail/client/config.go +++ b/pkg/promtail/client/config.go @@ -6,7 +6,7 @@ import ( "github.com/cortexproject/cortex/pkg/util" "github.com/cortexproject/cortex/pkg/util/flagext" - "github.com/prometheus/common/model" + lokiflag "github.com/grafana/loki/pkg/util/flagext" ) // Config describes configuration for a HTTP pusher client. @@ -17,8 +17,8 @@ type Config struct { BackoffConfig util.BackoffConfig `yaml:"backoff_config"` // The labels to add to any time series or alerts when communicating with loki - ExternalLabels model.LabelSet `yaml:"external_labels,omitempty"` - Timeout time.Duration `yaml:"timeout"` + ExternalLabels lokiflag.LabelSet `yaml:"external_labels,omitempty"` + Timeout time.Duration `yaml:"timeout"` } // RegisterFlags registers flags. @@ -31,7 +31,7 @@ func (c *Config) RegisterFlags(flags *flag.FlagSet) { flag.DurationVar(&c.BackoffConfig.MinBackoff, "client.min-backoff", 100*time.Millisecond, "Initial backoff time between retries.") flag.DurationVar(&c.BackoffConfig.MaxBackoff, "client.max-backoff", 5*time.Second, "Maximum backoff time between retries.") flag.DurationVar(&c.Timeout, "client.timeout", 10*time.Second, "Maximum time to wait for server to respond to a request") - + flags.Var(&c.ExternalLabels, "client.external-labels", "list of external labels to add to each log (e.g: --client.external-labels=lb1=v1,lb2=v2)") } // UnmarshalYAML implement Yaml Unmarshaler diff --git a/pkg/util/flagext/labelset.go b/pkg/util/flagext/labelset.go new file mode 100644 index 000000000000..c815694587e8 --- /dev/null +++ b/pkg/util/flagext/labelset.go @@ -0,0 +1,80 @@ +package flagext + +import ( + "bytes" + "encoding/csv" + "fmt" + "strings" + + "github.com/prometheus/common/model" +) + +// 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 +}