feat(canary): Support passing arbitrary set of labels to use for the query (#17008)

Signed-off-by: Noah Krause <krausenoah@gmail.com>
Co-authored-by: Ed Welch <ed@edjusted.com>

cmd/loki-canary/main.go

@@ -73,6 +73,7 @@ func main() {
 	buckets := flag.Int("buckets", 10, "Number of buckets in the response_latency histogram")
 
 	queryAppend := flag.String("query-append", "", "LogQL filters to be appended to the Canary query e.g. '| json | line_format `{{.log}}`'")
+	labels := flag.String("labels", "", "Comma-separated string of labels for the query e.g. 'service=loki,app=canary'. The parsing logic for this argument is simple, label values must not contain a comma or special characters and should not be quoted. Overwrites labelname and streamname")
 
 	metricTestInterval := flag.Duration("metric-test-interval", 1*time.Hour, "The interval the metric test query should be run")
 	metricTestQueryRange := flag.Duration("metric-test-range", 24*time.Hour, "The range value [24h] used in the metric test instant-query."+
@@ -207,7 +208,7 @@ func main() {
 
 		c.writer = writer.NewWriter(entryWriter, sentChan, *interval, *outOfOrderMin, *outOfOrderMax, *outOfOrderPercentage, *size, logger)
 		var err error
-		c.reader, err = reader.NewReader(os.Stderr, receivedChan, *useTLS, tlsConfig, *caFile, *certFile, *keyFile, *addr, *user, *pass, *tenantID, *queryTimeout, *lName, *lVal, *sName, *sValue, *interval, *queryAppend)
+		c.reader, err = reader.NewReader(os.Stderr, receivedChan, *useTLS, tlsConfig, *caFile, *certFile, *keyFile, *addr, *user, *pass, *tenantID, *queryTimeout, *lName, *lVal, *sName, *sValue, *interval, *queryAppend, *labels)
 		if err != nil {
 			_, _ = fmt.Fprintf(os.Stderr, "Unable to create reader for Loki querier, check config: %s", err)
 			os.Exit(1)

docs/sources/operations/loki-canary/_index.md

@@ -324,6 +324,8 @@ All options:
     	Duration between log entries (default 1s)
   -key-file string
     	Client PEM encoded X.509 key for optional use with TLS connection to Loki
+  -labels string
+        Comma-separated string of labels for the query e.g. 'service=loki,app=canary'. The parsing logic for this argument is simple, label values must not contain a comma or special characters and should not be quoted. Overwrites labelname and streamname
   -labelname string
     	The label name for this instance of loki-canary to use in the log selector (default "name")
   -labelvalue string
@@ -348,6 +350,8 @@ All options:
     	Frequency to check sent vs received logs, also the frequency which queries for missing logs will be dispatched to loki (default 1m0s)
   -push
     	Push the logs directly to given Loki address
+  -query-append string
+        LogQL filters to be appended to the Canary query e.g. '| json | line_format `{{.log}}`'  	
   -query-timeout duration
     	How long to wait for a query response from Loki (default 10s)
   -size int

pkg/canary/reader/reader.go

@@ -74,6 +74,7 @@ type Reader struct {
 	shuttingDown    bool
 	done            chan struct{}
 	queryAppend     string
+	labels          string
 }
 
 func NewReader(writer io.Writer,
@@ -92,6 +93,7 @@ func NewReader(writer io.Writer,
 	streamValue string,
 	interval time.Duration,
 	queryAppend string,
+	labels string,
 ) (*Reader, error) {
 	h := http.Header{}
 
@@ -154,6 +156,7 @@ func NewReader(writer io.Writer,
 		done:            make(chan struct{}),
 		shuttingDown:    false,
 		queryAppend:     queryAppend,
+		labels:          labels,
 	}
 
 	go rd.run()
@@ -290,12 +293,27 @@ func (r *Reader) Query(start time.Time, end time.Time) ([]time.Time, error) {
 	if r.useTLS {
 		scheme = "https"
 	}
+	var labels string
+	if r.labels != "" {
+		var lbls []string
+		for _, label := range strings.Split(r.labels, ",") {
+			labelParts := strings.Split(label, "=")
+			if len(labelParts) != 2 {
+				return nil, fmt.Errorf("invalid label format: %s, expected key=value", label)
+			}
+			lbls = append(lbls, fmt.Sprintf("%s=\"%s\"", labelParts[0], labelParts[1]))
+		}
+		labels = fmt.Sprintf("{%s}", strings.Join(lbls, ","))
+	} else {
+		labels = fmt.Sprintf("{%s=\"%s\",%s=\"%s\"}", r.sName, r.sValue, r.lName, r.lVal)
+	}
+
 	u := url.URL{
 		Scheme: scheme,
 		Host:   r.addr,
 		Path:   "/loki/api/v1/query_range",
 		RawQuery: fmt.Sprintf("start=%d&end=%d", start.UnixNano(), end.UnixNano()) +
-			"&query=" + url.QueryEscape(fmt.Sprintf("{%v=\"%v\",%v=\"%v\"} %v", r.sName, r.sValue, r.lName, r.lVal, r.queryAppend)) +
+			"&query=" + url.QueryEscape(fmt.Sprintf("%s %v", labels, r.queryAppend)) +
 			"&limit=1000",
 	}
 	fmt.Fprintf(r.w, "Querying loki for logs with query: %v\n", u.String())