adds external labels to be passed via flags (grafana#510)

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()

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

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
+}