move distributed query tenant propagation to protobuf to prevent header mismatch

Signed-off-by: Piyush Sharma <piyushsharma04321@gmail.com>

Author: ViciousEagle03

pkg/api/query/grpc.go

@@ -21,6 +21,7 @@ import (
 	"github.com/thanos-io/thanos/pkg/query"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb/prompb"
+	"github.com/thanos-io/thanos/pkg/tenancy"
 	"github.com/thanos-io/thanos/pkg/tracing"
 )
 
@@ -66,6 +67,10 @@ func RegisterQueryServer(queryServer querypb.QueryServer) func(*grpc.Server) {
 func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_QueryServer) error {
 	ctx := server.Context()
 
+	if request.Tenant != "" {
+		ctx = context.WithValue(ctx, tenancy.TenantKey, request.Tenant)
+	}
+
 	if request.TimeoutSeconds != 0 {
 		var cancel context.CancelFunc
 		timeout := time.Duration(request.TimeoutSeconds) * time.Second
@@ -189,6 +194,11 @@ func (g *GRPCAPI) Query(request *querypb.QueryRequest, server querypb.Query_Quer
 
 func (g *GRPCAPI) QueryRange(request *querypb.QueryRangeRequest, srv querypb.Query_QueryRangeServer) error {
 	ctx := srv.Context()
+
+	if request.Tenant != "" {
+		ctx = context.WithValue(ctx, tenancy.TenantKey, request.Tenant)
+	}
+
 	if request.TimeoutSeconds != 0 {
 		var cancel context.CancelFunc
 		timeout := time.Duration(request.TimeoutSeconds) * time.Second

pkg/api/query/querypb/query.pb.go

@@ -62,6 +62,8 @@ message QueryRequest {
   // If set to 0 or 1, each TimeSeries is sent as a separate response (default behavior).
   int64 response_batch_size = 15;
 
+  string tenant = 16;
+
   reserved 9;
 }
 
@@ -118,6 +120,8 @@ message QueryRangeRequest {
   // If set to 0 or 1, each TimeSeries is sent as a separate response (default behavior).
   int64 response_batch_size = 17;
 
+  string tenant = 18;
+  
   reserved 11;
 }
 

pkg/api/query/querypb/query.proto

@@ -22,7 +22,6 @@ import (
 	"github.com/prometheus/prometheus/util/annotations"
 	"github.com/prometheus/prometheus/util/stats"
 	"github.com/thanos-io/thanos/pkg/tenancy"
-	"google.golang.org/grpc/metadata"
 
 	"github.com/thanos-io/promql-engine/api"
 	"github.com/thanos-io/thanos/pkg/api/query/querypb"
@@ -340,9 +339,10 @@ func (r *remoteQuery) Exec(ctx context.Context) *promql.Result {
 
 	r.samplesStats = stats.NewQuerySamples(false)
 
+	var requestTenant string
 	if tenant := ctx.Value(tenancy.TenantKey); tenant != nil {
 		if tenantStr, ok := tenant.(string); ok {
-			qctx = metadata.AppendToOutgoingContext(qctx, tenancy.DefaultTenantHeader, tenantStr)
+			requestTenant = tenantStr
 		}
 	}
 
@@ -358,6 +358,7 @@ func (r *remoteQuery) Exec(ctx context.Context) *promql.Result {
 			MaxResolutionSeconds:  maxResolution,
 			EnableDedup:           true,
 			ResponseBatchSize:     store.DefaultResponseBatchSize,
+			Tenant:                requestTenant,
 		}
 
 		qry, err := r.client.Query(qctx, request)
@@ -437,6 +438,7 @@ func (r *remoteQuery) Exec(ctx context.Context) *promql.Result {
 		MaxResolutionSeconds:  maxResolution,
 		EnableDedup:           true,
 		ResponseBatchSize:     store.DefaultResponseBatchSize,
+		Tenant:                requestTenant,
 	}
 	qry, err := r.client.QueryRange(qctx, request)
 	if err != nil {