allow setting tenant id for querying logs from loki (#4185)

cmd/loki-canary/main.go

@@ -37,8 +37,9 @@ func main() {
 	port := flag.Int("port", 3500, "Port which loki-canary should expose metrics")
 	addr := flag.String("addr", "", "The Loki server URL:Port, e.g. loki:3100")
 	tls := flag.Bool("tls", false, "Does the loki connection use TLS?")
-	user := flag.String("user", "", "Loki username")
-	pass := flag.String("pass", "", "Loki password")
+	user := flag.String("user", "", "Loki username. Must not be set with tenant-id flag.")
+	pass := flag.String("pass", "", "Loki password. Must not be set with tenant-id flag.")
+	tenantID := flag.String("tenant-id", "", "Tenant id to be set in X-Scope-OrgID header. Must not be set with user/pass flags.")
 	queryTimeout := flag.Duration("query-timeout", 10*time.Second, "How long to wait for a query response from Loki")
 
 	interval := flag.Duration("interval", 1000*time.Millisecond, "Duration between log entries")
@@ -84,7 +85,7 @@ func main() {
 		defer c.lock.Unlock()
 
 		c.writer = writer.NewWriter(os.Stdout, sentChan, *interval, *size)
-		c.reader = reader.NewReader(os.Stderr, receivedChan, *tls, *addr, *user, *pass, *queryTimeout, *lName, *lVal, *sName, *sValue, *interval)
+		c.reader = reader.NewReader(os.Stderr, receivedChan, *tls, *addr, *user, *pass, *tenantID, *queryTimeout, *lName, *lVal, *sName, *sValue, *interval)
 		c.comparator = comparator.NewComparator(os.Stderr, *wait, *maxWait, *pruneInterval, *spotCheckInterval, *spotCheckMax, *spotCheckQueryRate, *spotCheckWait, *metricTestInterval, *metricTestQueryRange, *interval, *buckets, sentChan, receivedChan, c.reader, true)
 	}
 

pkg/canary/reader/reader.go

@@ -53,6 +53,7 @@ type Reader struct {
 	addr         string
 	user         string
 	pass         string
+	tenantID     string
 	queryTimeout time.Duration
 	sName        string
 	sValue       string
@@ -76,6 +77,7 @@ func NewReader(writer io.Writer,
 	address string,
 	user string,
 	pass string,
+	tenantID string,
 	queryTimeout time.Duration,
 	labelName string,
 	labelVal string,
@@ -85,6 +87,8 @@ func NewReader(writer io.Writer,
 	h := http.Header{}
 	if user != "" {
 		h = http.Header{"Authorization": {"Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass))}}
+	} else if tenantID != "" {
+		h = http.Header{"X-Scope-OrgID": {tenantID}}
 	}
 
 	next := time.Now()
@@ -101,6 +105,7 @@ func NewReader(writer io.Writer,
 		addr:         address,
 		user:         user,
 		pass:         pass,
+		tenantID:     tenantID,
 		queryTimeout: queryTimeout,
 		sName:        streamName,
 		sValue:       streamValue,
@@ -174,7 +179,11 @@ func (r *Reader) QueryCountOverTime(queryRange string) (float64, error) {
 		return 0, err
 	}
 
-	req.SetBasicAuth(r.user, r.pass)
+	if r.user != "" {
+		req.SetBasicAuth(r.user, r.pass)
+	} else if r.tenantID != "" {
+		req.Header.Set("X-Scope-OrgID", r.tenantID)
+	}
 	req.Header.Set("User-Agent", userAgent)
 
 	resp, err := http.DefaultClient.Do(req)
@@ -260,7 +269,11 @@ func (r *Reader) Query(start time.Time, end time.Time) ([]time.Time, error) {
 		return nil, err
 	}
 
-	req.SetBasicAuth(r.user, r.pass)
+	if r.user != "" {
+		req.SetBasicAuth(r.user, r.pass)
+	} else if r.tenantID != "" {
+		req.Header.Set("X-Scope-OrgID", r.tenantID)
+	}
 	req.Header.Set("User-Agent", userAgent)
 
 	resp, err := http.DefaultClient.Do(req)