diff --git a/pkg/storage/unified/resource/search_server_distributor.go b/pkg/storage/unified/resource/search_server_distributor.go index cceb0845b11..5fd4dacf7d9 100644 --- a/pkg/storage/unified/resource/search_server_distributor.go +++ b/pkg/storage/unified/resource/search_server_distributor.go @@ -1,11 +1,14 @@ package resource import ( + "cmp" "context" "errors" "fmt" "hash/fnv" "math/rand" + "slices" + "sort" "sync" "sync/atomic" "time" @@ -15,6 +18,7 @@ import ( "github.com/grafana/dskit/services" userutils "github.com/grafana/dskit/user" "github.com/prometheus/client_golang/prometheus" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" "google.golang.org/grpc" "google.golang.org/grpc/health/grpc_health_v1" @@ -35,10 +39,12 @@ func ProvideSearchDistributorServer(cfg *setting.Cfg, features featuremgmt.Featu } distributorServer := &distributorServer{ - log: log.New("index-server-distributor"), - ring: ring, - clientPool: ringClientPool, - tracing: tracer, + log: log.New("index-server-distributor"), + ring: ring, + clientPool: ringClientPool, + tracing: tracer, + subIndexesPerNamespace: cfg.SubIndexesPerNamespace, + replicationFactor: cfg.SearchRingReplicationFactor, } healthService, err := ProvideHealthService(distributorServer) @@ -83,10 +89,12 @@ const RingHeartbeatTimeout = time.Minute const RingNumTokens = 128 type distributorServer struct { - clientPool *ringclient.Pool - ring *ring.Ring - log log.Logger - tracing trace.Tracer + clientPool *ringclient.Pool + ring *ring.Ring + log log.Logger + tracing trace.Tracer + subIndexesPerNamespace int // Number of sub-indexes per namespace (0 = disabled) + replicationFactor int // Ring replication factor for replica failover } var ( @@ -99,12 +107,33 @@ var ( func (ds *distributorServer) Search(ctx context.Context, r *resourcepb.ResourceSearchRequest) (*resourcepb.ResourceSearchResponse, error) { ctx, span := ds.tracing.Start(ctx, "distributor.Search") defer span.End() - ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Options.Key.Namespace, "Search") - if err != nil { - return nil, err + + // If sub-index sharding is not enabled, use the existing single-node routing + if ds.subIndexesPerNamespace <= 0 { + ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Options.Key.Namespace, "Search") + if err != nil { + return nil, err + } + return client.Search(ctx, r) } - return client.Search(ctx, r) + // Scatter-gather search across all sub-indexes + nsr := NamespacedResource{ + Namespace: r.Options.Key.Namespace, + Group: r.Options.Key.Group, + Resource: r.Options.Key.Resource, + } + + span.SetAttributes( + attribute.String("namespace", nsr.Namespace), + attribute.String("group", nsr.Group), + attribute.String("resource", nsr.Resource), + attribute.Int("sub_indexes", ds.subIndexesPerNamespace), + ) + + subIndexes := ds.getSubIndexesForNamespace(nsr) + results := ds.parallelSearchWithFailover(ctx, subIndexes, r) + return ds.mergeResults(ctx, results, r) } func (ds *distributorServer) GetStats(ctx context.Context, r *resourcepb.ResourceStatsRequest) (*resourcepb.ResourceStatsResponse, error) { @@ -276,3 +305,383 @@ func (ds *distributorServer) IsHealthy(ctx context.Context, r *resourcepb.Health return &resourcepb.HealthCheckResponse{Status: resourcepb.HealthCheckResponse_NOT_SERVING}, nil } + +// --- Scatter-Gather Query Implementation with Replica Failover --- + +// subIndexSearchResult holds the result from searching a single sub-index. +type subIndexSearchResult struct { + subIndexID int + response *resourcepb.ResourceSearchResponse + err error + partialFailure bool // true if all replicas for this sub-index failed +} + +// getSubIndexesForNamespace returns all sub-index keys for a NamespacedResource. +// This is used for scatter-gather queries that need to query all sub-indexes. +func (ds *distributorServer) getSubIndexesForNamespace(nsr NamespacedResource) []SubIndexKey { + count := ds.subIndexesPerNamespace + if count <= 0 { + count = 1 + } + keys := make([]SubIndexKey, count) + for i := 0; i < count; i++ { + keys[i] = SubIndexKey{ + NamespacedResource: nsr, + SubIndexID: i, + } + } + return keys +} + +// getReplicasForSubIndex returns the ordered list of replicas (instances) that own +// the given sub-index. Replicas are ordered by preference from the ring. +// Uses consistent hashing including the sub-index ID to determine ownership. +func (ds *distributorServer) getReplicasForSubIndex(subIndex SubIndexKey) ([]ring.InstanceDesc, error) { + ringHasher := fnv.New32a() + // Include sub-index ID in hash to distribute sub-indexes across nodes + _, err := ringHasher.Write([]byte(fmt.Sprintf("%s/%d", subIndex.Namespace, subIndex.SubIndexID))) + if err != nil { + return nil, fmt.Errorf("error hashing sub-index key: %w", err) + } + + replicationFactor := ds.replicationFactor + if replicationFactor <= 0 { + replicationFactor = ds.ring.ReplicationFactor() + } + + rs, err := ds.ring.GetWithOptions(ringHasher.Sum32(), searchRingRead, ring.WithReplicationFactor(replicationFactor)) + if err != nil { + return nil, fmt.Errorf("error getting replication set from ring for sub-index %s: %w", subIndex.String(), err) + } + + return rs.Instances, nil +} + +// parallelSearchWithFailover executes search queries across all sub-indexes in parallel. +// For each sub-index, it tries replicas in order until one succeeds. +// Returns results from all sub-indexes (some may be marked as partial failures). +func (ds *distributorServer) parallelSearchWithFailover(ctx context.Context, subIndexes []SubIndexKey, r *resourcepb.ResourceSearchRequest) []*subIndexSearchResult { + ctx, span := ds.tracing.Start(ctx, "distributor.parallelSearchWithFailover") + defer span.End() + + results := make([]*subIndexSearchResult, len(subIndexes)) + var wg sync.WaitGroup + + for i, subIdx := range subIndexes { + wg.Add(1) + go func(idx int, key SubIndexKey) { + defer wg.Done() + results[idx] = ds.searchSubIndexWithFailover(ctx, key, r) + }(i, subIdx) + } + + wg.Wait() + + // Count partial failures for logging + partialFailures := 0 + for _, result := range results { + if result.partialFailure { + partialFailures++ + } + } + span.SetAttributes(attribute.Int("partial_failures", partialFailures)) + + return results +} + +// searchSubIndexWithFailover searches a single sub-index, trying each replica in order +// until one succeeds. If all replicas fail, marks the result as a partial failure. +func (ds *distributorServer) searchSubIndexWithFailover(ctx context.Context, subIndex SubIndexKey, r *resourcepb.ResourceSearchRequest) *subIndexSearchResult { + result := &subIndexSearchResult{ + subIndexID: subIndex.SubIndexID, + } + + // Get ordered list of replicas for this sub-index + replicas, err := ds.getReplicasForSubIndex(subIndex) + if err != nil { + ds.log.Warn("failed to get replicas for sub-index", "subIndex", subIndex.String(), "error", err) + result.err = err + result.partialFailure = true + return result + } + + if len(replicas) == 0 { + ds.log.Warn("no replicas available for sub-index", "subIndex", subIndex.String()) + result.err = fmt.Errorf("no replicas available for sub-index %s", subIndex.String()) + result.partialFailure = true + return result + } + + // Prepare context with metadata + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + md = make(metadata.MD) + } + rCtx := userutils.InjectOrgID(metadata.NewOutgoingContext(ctx, md), subIndex.Namespace) + + // Try each replica in order until success + var lastErr error + for replicaIdx, replica := range replicas { + client, err := ds.clientPool.GetClientForInstance(replica) + if err != nil { + ds.log.Debug("failed to get client for replica", + "subIndex", subIndex.String(), + "replica", replica.Id, + "replicaIdx", replicaIdx, + "error", err) + lastErr = err + continue + } + + resp, err := client.(*RingClient).Client.Search(rCtx, r) + if err == nil && resp.Error == nil { + // Success + result.response = resp + return result + } + + // Log failover event + if err != nil { + lastErr = err + ds.log.Warn("search failed on replica, failing over to next", + "subIndex", subIndex.String(), + "replica", replica.Id, + "replicaIdx", replicaIdx, + "remainingReplicas", len(replicas)-replicaIdx-1, + "error", err) + } else if resp.Error != nil { + lastErr = fmt.Errorf("search error: %s", resp.Error.Message) + ds.log.Warn("search returned error on replica, failing over to next", + "subIndex", subIndex.String(), + "replica", replica.Id, + "replicaIdx", replicaIdx, + "remainingReplicas", len(replicas)-replicaIdx-1, + "errorMessage", resp.Error.Message) + } + } + + // All replicas failed - mark as partial failure + ds.log.Error("all replicas failed for sub-index", + "subIndex", subIndex.String(), + "totalReplicas", len(replicas), + "lastError", lastErr) + result.err = lastErr + result.partialFailure = true + return result +} + +// mergeResults combines results from all sub-indexes into a single response. +// It handles: +// - Result deduplication by resource key +// - Sort merging (merge-sort for sorted results) +// - Pagination across shards +// - Tracking and reporting partial failures +func (ds *distributorServer) mergeResults(ctx context.Context, results []*subIndexSearchResult, req *resourcepb.ResourceSearchRequest) (*resourcepb.ResourceSearchResponse, error) { + ctx, span := ds.tracing.Start(ctx, "distributor.mergeResults") + defer span.End() + + // Collect all rows and track partial failures + var allRows []*resourcepb.ResourceTableRow + var columns []*resourcepb.ResourceTableColumnDefinition + var totalHits int64 + var totalQueryCost float64 + var maxScore float64 + partialFailures := 0 + failedSubIndexes := []int{} + facets := make(map[string]*resourcepb.ResourceSearchResponse_Facet) + + for _, result := range results { + if result.partialFailure { + partialFailures++ + failedSubIndexes = append(failedSubIndexes, result.subIndexID) + continue + } + + if result.response == nil { + continue + } + + resp := result.response + totalHits += resp.TotalHits + totalQueryCost += resp.QueryCost + if resp.MaxScore > maxScore { + maxScore = resp.MaxScore + } + + // Use columns from first valid response + if columns == nil && resp.Results != nil { + columns = resp.Results.Columns + } + + // Collect rows + if resp.Results != nil && len(resp.Results.Rows) > 0 { + allRows = append(allRows, resp.Results.Rows...) + } + + // Merge facets + for k, v := range resp.Facet { + if existing, ok := facets[k]; ok { + mergeFacets(existing, v) + } else { + facets[k] = v + } + } + } + + span.SetAttributes( + attribute.Int("total_rows_before_dedup", len(allRows)), + attribute.Int("partial_failures", partialFailures), + ) + + // Deduplicate rows by resource key + allRows = deduplicateRows(allRows) + + span.SetAttributes(attribute.Int("total_rows_after_dedup", len(allRows))) + + // Sort rows if sort criteria provided + if len(req.SortBy) > 0 && columns != nil { + sortRows(allRows, columns, req.SortBy) + } + + // Apply pagination + offset := int(req.Offset) + limit := int(req.Limit) + if limit <= 0 { + limit = 100 // default limit + } + + var paginatedRows []*resourcepb.ResourceTableRow + if offset < len(allRows) { + end := offset + limit + if end > len(allRows) { + end = len(allRows) + } + paginatedRows = allRows[offset:end] + } + + // Build response + response := &resourcepb.ResourceSearchResponse{ + TotalHits: totalHits, + QueryCost: totalQueryCost, + MaxScore: maxScore, + Results: &resourcepb.ResourceTable{ + Columns: columns, + Rows: paginatedRows, + }, + Facet: facets, + } + + // Report partial failures if any + if partialFailures > 0 { + response.Error = &resourcepb.ErrorResult{ + Code: 206, // Partial Content + Message: fmt.Sprintf("partial results: %d of %d sub-indexes failed (sub-indexes: %v)", partialFailures, len(results), failedSubIndexes), + } + ds.log.Warn("search returned partial results", + "namespace", req.Options.Key.Namespace, + "failedSubIndexes", partialFailures, + "totalSubIndexes", len(results)) + } + + return response, nil +} + +// deduplicateRows removes duplicate rows based on their resource key. +// If duplicates exist, keeps the first occurrence. +func deduplicateRows(rows []*resourcepb.ResourceTableRow) []*resourcepb.ResourceTableRow { + if len(rows) == 0 { + return rows + } + + seen := make(map[string]bool, len(rows)) + result := make([]*resourcepb.ResourceTableRow, 0, len(rows)) + + for _, row := range rows { + if row.Key == nil { + continue + } + key := SearchID(row.Key) + if !seen[key] { + seen[key] = true + result = append(result, row) + } + } + + return result +} + +// sortRows sorts rows based on the sort criteria. +// Uses stable sort to maintain relative ordering of equal elements. +func sortRows(rows []*resourcepb.ResourceTableRow, columns []*resourcepb.ResourceTableColumnDefinition, sortBy []*resourcepb.ResourceSearchRequest_Sort) { + if len(rows) == 0 || len(sortBy) == 0 { + return + } + + // Build column index map + columnIndex := make(map[string]int) + for i, col := range columns { + columnIndex[col.Name] = i + } + + sort.SliceStable(rows, func(i, j int) bool { + for _, s := range sortBy { + colIdx, ok := columnIndex[s.Field] + if !ok { + continue + } + + // Get cell values + var valI, valJ []byte + if colIdx < len(rows[i].Cells) { + valI = rows[i].Cells[colIdx] + } + if colIdx < len(rows[j].Cells) { + valJ = rows[j].Cells[colIdx] + } + + // Compare byte slices + cmpResult := slices.Compare(valI, valJ) + if cmpResult == 0 { + continue // Values are equal, check next sort field + } + + // Apply descending order if needed + if s.Desc { + return cmpResult > 0 + } + return cmpResult < 0 + } + return false // All sort fields are equal + }) +} + +// mergeFacets merges facet data from source into target. +func mergeFacets(target, source *resourcepb.ResourceSearchResponse_Facet) { + if source == nil { + return + } + + target.Total += source.Total + target.Missing += source.Missing + + // Merge term facets + termMap := make(map[string]int64) + for _, t := range target.Terms { + termMap[t.Term] = t.Count + } + for _, t := range source.Terms { + termMap[t.Term] += t.Count + } + + // Rebuild term slice sorted by count (descending) + target.Terms = make([]*resourcepb.ResourceSearchResponse_TermFacet, 0, len(termMap)) + for term, count := range termMap { + target.Terms = append(target.Terms, &resourcepb.ResourceSearchResponse_TermFacet{ + Term: term, + Count: count, + }) + } + slices.SortFunc(target.Terms, func(a, b *resourcepb.ResourceSearchResponse_TermFacet) int { + return cmp.Compare(b.Count, a.Count) // Descending order + }) +} diff --git a/pkg/storage/unified/resource/search_server_distributor_test.go b/pkg/storage/unified/resource/search_server_distributor_test.go new file mode 100644 index 00000000000..4c3b6929443 --- /dev/null +++ b/pkg/storage/unified/resource/search_server_distributor_test.go @@ -0,0 +1,390 @@ +package resource + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/trace/noop" + + "github.com/grafana/grafana/pkg/infra/log" + "github.com/grafana/grafana/pkg/storage/unified/resourcepb" +) + +func TestGetSubIndexesForNamespace(t *testing.T) { + tests := []struct { + name string + subIndexCount int + expectedCount int + expectedIDs []int + }{ + { + name: "zero sub-indexes defaults to 1", + subIndexCount: 0, + expectedCount: 1, + expectedIDs: []int{0}, + }, + { + name: "negative sub-indexes defaults to 1", + subIndexCount: -1, + expectedCount: 1, + expectedIDs: []int{0}, + }, + { + name: "4 sub-indexes", + subIndexCount: 4, + expectedCount: 4, + expectedIDs: []int{0, 1, 2, 3}, + }, + { + name: "64 sub-indexes", + subIndexCount: 64, + expectedCount: 64, + expectedIDs: nil, // Don't check all 64 + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ds := &distributorServer{ + subIndexesPerNamespace: tt.subIndexCount, + } + + nsr := NamespacedResource{ + Namespace: "org-1", + Group: "dashboard.grafana.app", + Resource: "dashboards", + } + + result := ds.getSubIndexesForNamespace(nsr) + assert.Len(t, result, tt.expectedCount) + + // Check that all keys have the correct NSR + for i, key := range result { + assert.Equal(t, nsr, key.NamespacedResource) + assert.Equal(t, i, key.SubIndexID) + } + + // Check specific IDs if provided + if tt.expectedIDs != nil { + for i, expectedID := range tt.expectedIDs { + assert.Equal(t, expectedID, result[i].SubIndexID) + } + } + }) + } +} + +func TestDeduplicateRows(t *testing.T) { + tests := []struct { + name string + input []*resourcepb.ResourceTableRow + expected int // expected number of rows after dedup + }{ + { + name: "empty input", + input: []*resourcepb.ResourceTableRow{}, + expected: 0, + }, + { + name: "no duplicates", + input: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "c"}}, + }, + expected: 3, + }, + { + name: "with duplicates", + input: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}}, // duplicate + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "c"}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}}, // duplicate + }, + expected: 3, + }, + { + name: "rows with nil keys are skipped", + input: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}}, + {Key: nil}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}}, + }, + expected: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := deduplicateRows(tt.input) + assert.Len(t, result, tt.expected) + }) + } +} + +func TestSortRows(t *testing.T) { + columns := []*resourcepb.ResourceTableColumnDefinition{ + {Name: "title"}, + {Name: "created"}, + } + + tests := []struct { + name string + rows []*resourcepb.ResourceTableRow + sortBy []*resourcepb.ResourceSearchRequest_Sort + expectedOrder []string // expected order of first cell values + }{ + { + name: "sort ascending by title", + rows: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Name: "c"}, Cells: [][]byte{[]byte("charlie"), nil}}, + {Key: &resourcepb.ResourceKey{Name: "a"}, Cells: [][]byte{[]byte("alpha"), nil}}, + {Key: &resourcepb.ResourceKey{Name: "b"}, Cells: [][]byte{[]byte("bravo"), nil}}, + }, + sortBy: []*resourcepb.ResourceSearchRequest_Sort{ + {Field: "title", Desc: false}, + }, + expectedOrder: []string{"alpha", "bravo", "charlie"}, + }, + { + name: "sort descending by title", + rows: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Name: "a"}, Cells: [][]byte{[]byte("alpha"), nil}}, + {Key: &resourcepb.ResourceKey{Name: "b"}, Cells: [][]byte{[]byte("bravo"), nil}}, + {Key: &resourcepb.ResourceKey{Name: "c"}, Cells: [][]byte{[]byte("charlie"), nil}}, + }, + sortBy: []*resourcepb.ResourceSearchRequest_Sort{ + {Field: "title", Desc: true}, + }, + expectedOrder: []string{"charlie", "bravo", "alpha"}, + }, + { + name: "empty rows", + rows: []*resourcepb.ResourceTableRow{}, + sortBy: []*resourcepb.ResourceSearchRequest_Sort{{Field: "title"}}, + expectedOrder: []string{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sortRows(tt.rows, columns, tt.sortBy) + + for i, expected := range tt.expectedOrder { + if i < len(tt.rows) { + assert.Equal(t, expected, string(tt.rows[i].Cells[0])) + } + } + }) + } +} + +func TestMergeFacets(t *testing.T) { + t.Run("merge term facets", func(t *testing.T) { + target := &resourcepb.ResourceSearchResponse_Facet{ + Field: "tags", + Total: 100, + Missing: 5, + Terms: []*resourcepb.ResourceSearchResponse_TermFacet{ + {Term: "production", Count: 50}, + {Term: "staging", Count: 30}, + }, + } + + source := &resourcepb.ResourceSearchResponse_Facet{ + Field: "tags", + Total: 80, + Missing: 3, + Terms: []*resourcepb.ResourceSearchResponse_TermFacet{ + {Term: "production", Count: 40}, + {Term: "development", Count: 25}, + }, + } + + mergeFacets(target, source) + + assert.Equal(t, int64(180), target.Total) + assert.Equal(t, int64(8), target.Missing) + assert.Len(t, target.Terms, 3) + + // Terms should be sorted by count descending + termCounts := make(map[string]int64) + for _, term := range target.Terms { + termCounts[term.Term] = term.Count + } + assert.Equal(t, int64(90), termCounts["production"]) + assert.Equal(t, int64(30), termCounts["staging"]) + assert.Equal(t, int64(25), termCounts["development"]) + }) + + t.Run("merge nil source", func(t *testing.T) { + target := &resourcepb.ResourceSearchResponse_Facet{ + Total: 100, + } + mergeFacets(target, nil) + assert.Equal(t, int64(100), target.Total) + }) +} + +func TestSubIndexSearchResult(t *testing.T) { + t.Run("successful result", func(t *testing.T) { + result := &subIndexSearchResult{ + subIndexID: 5, + response: &resourcepb.ResourceSearchResponse{ + TotalHits: 100, + }, + partialFailure: false, + } + assert.False(t, result.partialFailure) + assert.NotNil(t, result.response) + }) + + t.Run("partial failure result", func(t *testing.T) { + result := &subIndexSearchResult{ + subIndexID: 3, + err: assert.AnError, + partialFailure: true, + } + assert.True(t, result.partialFailure) + assert.Nil(t, result.response) + assert.Error(t, result.err) + }) +} + +func TestSubIndexKey(t *testing.T) { + key := SubIndexKey{ + NamespacedResource: NamespacedResource{ + Namespace: "org-1", + Group: "dashboard.grafana.app", + Resource: "dashboards", + }, + SubIndexID: 42, + } + + t.Run("String representation", func(t *testing.T) { + expected := "org-1/dashboard.grafana.app/dashboards/shard-42" + assert.Equal(t, expected, key.String()) + }) + + t.Run("ToNamespacedResource", func(t *testing.T) { + nsr := key.ToNamespacedResource() + assert.Equal(t, "org-1", nsr.Namespace) + assert.Equal(t, "dashboard.grafana.app", nsr.Group) + assert.Equal(t, "dashboards", nsr.Resource) + }) +} + +func newTestDistributorServer() *distributorServer { + return &distributorServer{ + tracing: noop.NewTracerProvider().Tracer("test"), + log: log.New("test-distributor"), + } +} + +func TestMergeResultsPagination(t *testing.T) { + // Create mock results from 3 sub-indexes + results := []*subIndexSearchResult{ + { + subIndexID: 0, + response: &resourcepb.ResourceSearchResponse{ + TotalHits: 10, + Results: &resourcepb.ResourceTable{ + Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}}, + Rows: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a1"}, Cells: [][]byte{[]byte("a1")}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a2"}, Cells: [][]byte{[]byte("a2")}}, + }, + }, + }, + }, + { + subIndexID: 1, + response: &resourcepb.ResourceSearchResponse{ + TotalHits: 10, + Results: &resourcepb.ResourceTable{ + Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}}, + Rows: []*resourcepb.ResourceTableRow{ + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b1"}, Cells: [][]byte{[]byte("b1")}}, + {Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b2"}, Cells: [][]byte{[]byte("b2")}}, + }, + }, + }, + }, + } + + t.Run("pagination with limit", func(t *testing.T) { + ds := newTestDistributorServer() + req := &resourcepb.ResourceSearchRequest{ + Limit: 2, + Offset: 0, + Options: &resourcepb.ListOptions{ + Key: &resourcepb.ResourceKey{Namespace: "ns"}, + }, + } + + resp, err := ds.mergeResults(context.Background(), results, req) + require.NoError(t, err) + assert.Equal(t, int64(20), resp.TotalHits) // 10 + 10 + assert.Len(t, resp.Results.Rows, 2) + }) + + t.Run("pagination with offset", func(t *testing.T) { + ds := newTestDistributorServer() + req := &resourcepb.ResourceSearchRequest{ + Limit: 2, + Offset: 2, + Options: &resourcepb.ListOptions{ + Key: &resourcepb.ResourceKey{Namespace: "ns"}, + }, + } + + resp, err := ds.mergeResults(context.Background(), results, req) + require.NoError(t, err) + assert.Len(t, resp.Results.Rows, 2) // rows 3 and 4 + }) +} + +func TestMergeResultsPartialFailure(t *testing.T) { + results := []*subIndexSearchResult{ + { + subIndexID: 0, + response: &resourcepb.ResourceSearchResponse{ + TotalHits: 10, + Results: &resourcepb.ResourceTable{ + Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}}, + Rows: []*resourcepb.ResourceTableRow{}, + }, + }, + }, + { + subIndexID: 1, + partialFailure: true, + err: assert.AnError, + }, + { + subIndexID: 2, + partialFailure: true, + err: assert.AnError, + }, + } + + ds := newTestDistributorServer() + req := &resourcepb.ResourceSearchRequest{ + Limit: 100, + Options: &resourcepb.ListOptions{ + Key: &resourcepb.ResourceKey{Namespace: "ns"}, + }, + } + + resp, err := ds.mergeResults(context.Background(), results, req) + require.NoError(t, err) + + // Should have partial failure error + require.NotNil(t, resp.Error) + assert.Equal(t, int32(206), resp.Error.Code) // HTTP 206 Partial Content + assert.Contains(t, resp.Error.Message, "2 of 3 sub-indexes failed") +}