From 653d9fa895907363d5ed3a544e21cf7dddb8e5d1 Mon Sep 17 00:00:00 2001
From: Fabrizio <135109076+fabrizio-grafana@users.noreply.github.com>
Date: Thu, 11 Apr 2024 11:51:54 +0200
Subject: [PATCH] Loki: Handle `X-Scope-OrgID` and tenant IDs (#85726)

---
 pkg/tsdb/loki/api.go | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/pkg/tsdb/loki/api.go b/pkg/tsdb/loki/api.go
index 6ffa5226a65..f914063f095 100644
--- a/pkg/tsdb/loki/api.go
+++ b/pkg/tsdb/loki/api.go
@@ -18,9 +18,9 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/promlib/converter"
@@ -109,6 +109,8 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery, cat
 		req.Header.Set("X-Loki-Response-Encoding-Flags", "categorize-labels")
 	}
 
+	setXScopeOrgIDHeader(req, ctx)
+
 	return req, nil
 }
 
@@ -257,6 +259,8 @@ func makeRawRequest(ctx context.Context, lokiDsUrl string, resourcePath string)
 		return nil, err
 	}
 
+	setXScopeOrgIDHeader(req, ctx)
+
 	return req, nil
 }
 
@@ -346,3 +350,29 @@ func getSupportingQueryHeaderValue(supportingQueryType SupportingQueryType) stri
 
 	return value
 }
+
+// setXScopeOrgIDHeader sets the `X-Scope-OrgID` header in the provided HTTP request based on the tenant ID retrieved from the context.
+// `X-Scope-OrgID` is needed by the Loki system to work in multi-tenant mode.
+// See https://github.com/grafana/loki/blob/main/docs/sources/operations/multi-tenancy.md
+func setXScopeOrgIDHeader(req *http.Request, ctx context.Context) *http.Request {
+	logger := backend.NewLoggerWith("logger", "tsdb.loki")
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		logger.Error("Error in retrieving metadata from context. Header not set")
+		return req
+	}
+
+	tenantids := md.Get("tenantid")
+	if len(tenantids) == 0 {
+		// We assume we are not using multi-tenant mode, which is fine
+		logger.Debug("Tenant ID not present. Header not set")
+	} else if len(tenantids[0]) > 1 {
+		// Loki supports multiple tenant IDs, but we should receive them from different contexts
+		logger.Error(strconv.Itoa(len(tenantids)) + " tenant IDs found. Header not set")
+	} else {
+		req.Header.Add("X-Scope-OrgID", tenantids[0])
+		logger.Debug("Tenant ID " + tenantids[0] + " added to Loki request")
+	}
+	return req
+}