Merge pull request #18569 from roidelapluie/roidelapluie/labelnames-limit

tsdb: apply LabelNames limit from LabelHints in blockBaseQuerier

tsdb/querier.go

@@ -81,9 +81,17 @@ func (q *blockBaseQuerier) LabelValues(ctx context.Context, name string, hints *
 	return res, nil, err
 }
 
-func (q *blockBaseQuerier) LabelNames(ctx context.Context, _ *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
+func (q *blockBaseQuerier) LabelNames(ctx context.Context, hints *storage.LabelHints, matchers ...*labels.Matcher) ([]string, annotations.Annotations, error) {
 	res, err := q.index.LabelNames(ctx, matchers...)
-	return res, nil, err
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if hints != nil && hints.Limit > 0 && len(res) > hints.Limit {
+		res = res[:hints.Limit]
+	}
+
+	return res, nil, nil
 }
 
 func (q *blockBaseQuerier) Close() error {

tsdb/querier_test.go

@@ -418,6 +418,41 @@ func TestBlockQuerier(t *testing.T) {
 	}
 }
 
+func TestBlockBaseQuerier_LabelNamesLimit(t *testing.T) {
+	ix := newMockIndex()
+	ls := labels.FromStrings("aaa", "1", "bbb", "2", "ccc", "3")
+	require.NoError(t, ix.AddSeries(1, ls))
+	ls.Range(func(lbl labels.Label) {
+		require.NoError(t, ix.WritePostings(lbl.Name, lbl.Value, index.NewListPostings([]storage.SeriesRef{1})))
+	})
+
+	q := &blockBaseQuerier{index: ix, chunks: mockChunkReader(nil), tombstones: tombstones.NewMemTombstones()}
+
+	t.Run("no limit", func(t *testing.T) {
+		names, _, err := q.LabelNames(context.Background(), nil)
+		require.NoError(t, err)
+		require.Equal(t, []string{"aaa", "bbb", "ccc"}, names)
+	})
+
+	t.Run("limit applied", func(t *testing.T) {
+		names, _, err := q.LabelNames(context.Background(), &storage.LabelHints{Limit: 2})
+		require.NoError(t, err)
+		require.Equal(t, []string{"aaa", "bbb"}, names)
+	})
+
+	t.Run("limit larger than result", func(t *testing.T) {
+		names, _, err := q.LabelNames(context.Background(), &storage.LabelHints{Limit: 10})
+		require.NoError(t, err)
+		require.Equal(t, []string{"aaa", "bbb", "ccc"}, names)
+	})
+
+	t.Run("zero limit means no limit", func(t *testing.T) {
+		names, _, err := q.LabelNames(context.Background(), &storage.LabelHints{Limit: 0})
+		require.NoError(t, err)
+		require.Equal(t, []string{"aaa", "bbb", "ccc"}, names)
+	})
+}
+
 func TestBlockQuerier_AgainstHeadWithOpenChunks(t *testing.T) {
 	for _, c := range []blockQuerierTestCase{
 		{