From e2be9f04a5f41da53c717cf6328fdec23a2a1d99 Mon Sep 17 00:00:00 2001
From: Charlie Le <charlie_le@apple.com>
Date: Fri, 10 May 2024 08:40:56 -0700
Subject: [PATCH] Select correct tenant during query federation

We are creating a job which has the tenant name for each tenant to
query, but we use the incorrect index in the `ids` slice that contains
the list of tenants. This causes the querier to select results for the
first tenant in the list of tenants to match with instead of the actual
tenant the querier was intended for.

I updated the existing tests to check the labels of the results as well
instead of only relying on the number of series found.

Fixes https://github.com/cortexproject/cortex/issues/5941

Signed-off-by: Charlie Le <charlie_le@apple.com>
---
 .../tenantfederation/merge_queryable.go       |  2 +-
 .../tenantfederation/merge_queryable_test.go  | 44 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/pkg/querier/tenantfederation/merge_queryable.go b/pkg/querier/tenantfederation/merge_queryable.go
index c765f9f2d3..c172261d1e 100644
--- a/pkg/querier/tenantfederation/merge_queryable.go
+++ b/pkg/querier/tenantfederation/merge_queryable.go
@@ -339,7 +339,7 @@ func (m *mergeQuerier) Select(ctx context.Context, sortSeries bool, hints *stora
 			return fmt.Errorf("unexpected type %T", jobIntf)
 		}
 		// Based on parent ctx here as we are using lazy querier.
-		newCtx := user.InjectOrgID(parentCtx, ids[job.pos])
+		newCtx := user.InjectOrgID(parentCtx, job.id)
 		seriesSets[job.pos] = &addLabelsSeriesSet{
 			upstream: job.querier.Select(newCtx, sortSeries, hints, filteredMatchers...),
 			labels: labels.Labels{
diff --git a/pkg/querier/tenantfederation/merge_queryable_test.go b/pkg/querier/tenantfederation/merge_queryable_test.go
index e8aa04ea28..f0b985a0a7 100644
--- a/pkg/querier/tenantfederation/merge_queryable_test.go
+++ b/pkg/querier/tenantfederation/merge_queryable_test.go
@@ -446,11 +446,42 @@ func TestMergeQueryable_Select(t *testing.T) {
 					name:                "should return only series for team-a and team-c tenants when there is a not-equals matcher for the team-b tenant",
 					matchers:            []*labels.Matcher{{Name: defaultTenantLabel, Value: "team-b", Type: labels.MatchNotEqual}},
 					expectedSeriesCount: 4,
+					expectedLabels: []labels.Labels{
+						{
+							{Name: "__tenant_id__", Value: "team-a"},
+							{Name: "instance", Value: "host1"},
+							{Name: "tenant-team-a", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-a"},
+							{Name: "instance", Value: "host2.team-a"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-c"},
+							{Name: "instance", Value: "host1"},
+							{Name: "tenant-team-c", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-c"},
+							{Name: "instance", Value: "host2.team-c"},
+						},
+					},
 				},
 				{
 					name:                "should return only series for team-b when there is an equals matcher for the team-b tenant",
 					matchers:            []*labels.Matcher{{Name: defaultTenantLabel, Value: "team-b", Type: labels.MatchEqual}},
 					expectedSeriesCount: 2,
+					expectedLabels: []labels.Labels{
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host1"},
+							{Name: "tenant-team-b", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host2.team-b"},
+						},
+					},
 				},
 				{
 					name:                "should return one series for each tenant when there is an equals matcher for the host1 instance",
@@ -516,6 +547,19 @@ func TestMergeQueryable_Select(t *testing.T) {
 					name:                "should return only series for team-b when there is an equals matcher for team-b tenant",
 					matchers:            []*labels.Matcher{{Name: defaultTenantLabel, Value: "team-b", Type: labels.MatchEqual}},
 					expectedSeriesCount: 2,
+					expectedLabels: []labels.Labels{
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host1"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+							{Name: "tenant-team-b", Value: "static"},
+						},
+						{
+							{Name: "__tenant_id__", Value: "team-b"},
+							{Name: "instance", Value: "host2.team-b"},
+							{Name: "original___tenant_id__", Value: "original-value"},
+						},
+					},
 				},
 				{
 					name:                "should return all series when there is an equals matcher for the original value of __tenant_id__ using the revised tenant label",