Add context argument to LabelValuesStream

Signed-off-by: Arve Knudsen <[email protected]>
grafana · Oct 5, 2023 · c03952d · c03952d
1 parent c367f6a
commit c03952d
Show file tree

Hide file tree

Showing 23 changed files with 128 additions and 120 deletions.
diff --git a/model/labels/regexp_test.go b/model/labels/regexp_test.go
@@ -27,7 +27,6 @@ import (
 	"github.com/DmitriyVTitov/size"
 	"github.com/grafana/regexp"
 	"github.com/grafana/regexp/syntax"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -511,17 +510,17 @@ func TestStringMatcherFromRegexp_LiteralPrefix(t *testing.T) {
 			require.Equal(t, c.expectedLiteralPrefixMatchers, numPrefixMatchers)
 
 			for _, value := range c.expectedMatches {
-				assert.Truef(t, matcher.Matches(value), "Value: %s", value)
+				require.Truef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Truef(t, re.MatchString(value), "Value: %s", value)
+				require.Truef(t, re.MatchString(value), "Value: %s", value)
 			}
 
 			for _, value := range c.expectedNotMatches {
-				assert.Falsef(t, matcher.Matches(value), "Value: %s", value)
+				require.Falsef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Falsef(t, re.MatchString(value), "Value: %s", value)
+				require.Falsef(t, re.MatchString(value), "Value: %s", value)
 			}
 		})
 	}
@@ -586,17 +585,17 @@ func TestStringMatcherFromRegexp_LiteralSuffix(t *testing.T) {
 			require.Equal(t, c.expectedLiteralSuffixMatchers, numSuffixMatchers)
 
 			for _, value := range c.expectedMatches {
-				assert.Truef(t, matcher.Matches(value), "Value: %s", value)
+				require.Truef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Truef(t, re.MatchString(value), "Value: %s", value)
+				require.Truef(t, re.MatchString(value), "Value: %s", value)
 			}
 
 			for _, value := range c.expectedNotMatches {
-				assert.Falsef(t, matcher.Matches(value), "Value: %s", value)
+				require.Falsef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Falsef(t, re.MatchString(value), "Value: %s", value)
+				require.Falsef(t, re.MatchString(value), "Value: %s", value)
 			}
 		})
 	}
@@ -671,17 +670,17 @@ func TestStringMatcherFromRegexp_Quest(t *testing.T) {
 			require.Equal(t, c.expectedZeroOrOneMatchers, numZeroOrOneMatchers)
 
 			for _, value := range c.expectedMatches {
-				assert.Truef(t, matcher.Matches(value), "Value: %s", value)
+				require.Truef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Truef(t, re.MatchString(value), "Value: %s", value)
+				require.Truef(t, re.MatchString(value), "Value: %s", value)
 			}
 
 			for _, value := range c.expectedNotMatches {
-				assert.Falsef(t, matcher.Matches(value), "Value: %s", value)
+				require.Falsef(t, matcher.Matches(value), "Value: %s", value)
 
 				// Ensure the golang regexp engine would return the same.
-				assert.Falsef(t, re.MatchString(value), "Value: %s", value)
+				require.Falsef(t, re.MatchString(value), "Value: %s", value)
 			}
 		})
 	}
@@ -1116,8 +1115,8 @@ func TestFindEqualStringMatchers(t *testing.T) {
 
 	t.Run("empty matcher", func(t *testing.T) {
 		actualMatches, actualOk := findEqualStringMatchersAndCollectMatches(emptyStringMatcher{})
-		assert.False(t, actualOk)
-		assert.Empty(t, actualMatches)
+		require.False(t, actualOk)
+		require.Empty(t, actualMatches)
 	})
 
 	t.Run("concat of literal matchers (case sensitive)", func(t *testing.T) {
@@ -1128,8 +1127,8 @@ func TestFindEqualStringMatchers(t *testing.T) {
 			},
 		)
 
-		assert.True(t, actualOk)
-		assert.Equal(t, []match{{"test-1", true}, {"test-2", true}}, actualMatches)
+		require.True(t, actualOk)
+		require.Equal(t, []match{{"test-1", true}, {"test-2", true}}, actualMatches)
 	})
 
 	t.Run("concat of literal matchers (case insensitive)", func(t *testing.T) {
@@ -1140,8 +1139,8 @@ func TestFindEqualStringMatchers(t *testing.T) {
 			},
 		)
 
-		assert.True(t, actualOk)
-		assert.Equal(t, []match{{"test-1", false}, {"test-2", false}}, actualMatches)
+		require.True(t, actualOk)
+		require.Equal(t, []match{{"test-1", false}, {"test-2", false}}, actualMatches)
 	})
 
 	t.Run("concat of literal matchers (mixed case)", func(t *testing.T) {
@@ -1152,8 +1151,8 @@ func TestFindEqualStringMatchers(t *testing.T) {
 			},
 		)
 
-		assert.True(t, actualOk)
-		assert.Equal(t, []match{{"test-1", false}, {"test-2", true}}, actualMatches)
+		require.True(t, actualOk)
+		require.Equal(t, []match{{"test-1", false}, {"test-2", true}}, actualMatches)
 	})
 }
 
@@ -1211,86 +1210,86 @@ func BenchmarkOptimizeEqualStringMatchers(b *testing.B) {
 
 func TestZeroOrOneCharacterStringMatcher(t *testing.T) {
 	matcher := &zeroOrOneCharacterStringMatcher{matchNL: true}
-	assert.True(t, matcher.Matches(""))
-	assert.True(t, matcher.Matches("x"))
-	assert.True(t, matcher.Matches("\n"))
-	assert.False(t, matcher.Matches("xx"))
-	assert.False(t, matcher.Matches("\n\n"))
+	require.True(t, matcher.Matches(""))
+	require.True(t, matcher.Matches("x"))
+	require.True(t, matcher.Matches("\n"))
+	require.False(t, matcher.Matches("xx"))
+	require.False(t, matcher.Matches("\n\n"))
 
 	matcher = &zeroOrOneCharacterStringMatcher{matchNL: false}
-	assert.True(t, matcher.Matches(""))
-	assert.True(t, matcher.Matches("x"))
-	assert.False(t, matcher.Matches("\n"))
-	assert.False(t, matcher.Matches("xx"))
-	assert.False(t, matcher.Matches("\n\n"))
+	require.True(t, matcher.Matches(""))
+	require.True(t, matcher.Matches("x"))
+	require.False(t, matcher.Matches("\n"))
+	require.False(t, matcher.Matches("xx"))
+	require.False(t, matcher.Matches("\n\n"))
 }
 
 func TestLiteralPrefixStringMatcher(t *testing.T) {
 	m := &literalPrefixStringMatcher{prefix: "mar", prefixCaseSensitive: true, right: &emptyStringMatcher{}}
-	assert.True(t, m.Matches("mar"))
-	assert.False(t, m.Matches("marco"))
-	assert.False(t, m.Matches("ma"))
-	assert.False(t, m.Matches("mAr"))
+	require.True(t, m.Matches("mar"))
+	require.False(t, m.Matches("marco"))
+	require.False(t, m.Matches("ma"))
+	require.False(t, m.Matches("mAr"))
 
 	m = &literalPrefixStringMatcher{prefix: "mar", prefixCaseSensitive: false, right: &emptyStringMatcher{}}
-	assert.True(t, m.Matches("mar"))
-	assert.False(t, m.Matches("marco"))
-	assert.False(t, m.Matches("ma"))
-	assert.True(t, m.Matches("mAr"))
+	require.True(t, m.Matches("mar"))
+	require.False(t, m.Matches("marco"))
+	require.False(t, m.Matches("ma"))
+	require.True(t, m.Matches("mAr"))
 
 	m = &literalPrefixStringMatcher{prefix: "mar", prefixCaseSensitive: true, right: &equalStringMatcher{s: "co", caseSensitive: false}}
-	assert.True(t, m.Matches("marco"))
-	assert.True(t, m.Matches("marCO"))
-	assert.False(t, m.Matches("MARco"))
-	assert.False(t, m.Matches("mar"))
-	assert.False(t, m.Matches("marcopracucci"))
+	require.True(t, m.Matches("marco"))
+	require.True(t, m.Matches("marCO"))
+	require.False(t, m.Matches("MARco"))
+	require.False(t, m.Matches("mar"))
+	require.False(t, m.Matches("marcopracucci"))
 }
 
 func TestLiteralSuffixStringMatcher(t *testing.T) {
 	m := &literalSuffixStringMatcher{left: &emptyStringMatcher{}, suffix: "co", suffixCaseSensitive: true}
-	assert.True(t, m.Matches("co"))
-	assert.False(t, m.Matches("marco"))
-	assert.False(t, m.Matches("coo"))
-	assert.False(t, m.Matches("Co"))
+	require.True(t, m.Matches("co"))
+	require.False(t, m.Matches("marco"))
+	require.False(t, m.Matches("coo"))
+	require.False(t, m.Matches("Co"))
 
 	m = &literalSuffixStringMatcher{left: &emptyStringMatcher{}, suffix: "co", suffixCaseSensitive: false}
-	assert.True(t, m.Matches("co"))
-	assert.False(t, m.Matches("marco"))
-	assert.False(t, m.Matches("coo"))
-	assert.True(t, m.Matches("Co"))
+	require.True(t, m.Matches("co"))
+	require.False(t, m.Matches("marco"))
+	require.False(t, m.Matches("coo"))
+	require.True(t, m.Matches("Co"))
 
 	m = &literalSuffixStringMatcher{left: &equalStringMatcher{s: "mar", caseSensitive: false}, suffix: "co", suffixCaseSensitive: true}
-	assert.True(t, m.Matches("marco"))
-	assert.True(t, m.Matches("MARco"))
-	assert.False(t, m.Matches("marCO"))
-	assert.False(t, m.Matches("mar"))
-	assert.False(t, m.Matches("marcopracucci"))
+	require.True(t, m.Matches("marco"))
+	require.True(t, m.Matches("MARco"))
+	require.False(t, m.Matches("marCO"))
+	require.False(t, m.Matches("mar"))
+	require.False(t, m.Matches("marcopracucci"))
 
 	m = &literalSuffixStringMatcher{left: &equalStringMatcher{s: "mar", caseSensitive: false}, suffix: "co", suffixCaseSensitive: false}
-	assert.True(t, m.Matches("marco"))
-	assert.True(t, m.Matches("MARco"))
-	assert.True(t, m.Matches("marCO"))
-	assert.False(t, m.Matches("mar"))
-	assert.False(t, m.Matches("marcopracucci"))
+	require.True(t, m.Matches("marco"))
+	require.True(t, m.Matches("MARco"))
+	require.True(t, m.Matches("marCO"))
+	require.False(t, m.Matches("mar"))
+	require.False(t, m.Matches("marcopracucci"))
 }
 
 func TestHasPrefixCaseInsensitive(t *testing.T) {
-	assert.True(t, hasPrefixCaseInsensitive("marco", "mar"))
-	assert.True(t, hasPrefixCaseInsensitive("mArco", "mar"))
-	assert.True(t, hasPrefixCaseInsensitive("marco", "MaR"))
-	assert.True(t, hasPrefixCaseInsensitive("marco", "marco"))
-	assert.True(t, hasPrefixCaseInsensitive("mArco", "marco"))
-
-	assert.False(t, hasPrefixCaseInsensitive("marco", "a"))
-	assert.False(t, hasPrefixCaseInsensitive("marco", "abcdefghi"))
+	require.True(t, hasPrefixCaseInsensitive("marco", "mar"))
+	require.True(t, hasPrefixCaseInsensitive("mArco", "mar"))
+	require.True(t, hasPrefixCaseInsensitive("marco", "MaR"))
+	require.True(t, hasPrefixCaseInsensitive("marco", "marco"))
+	require.True(t, hasPrefixCaseInsensitive("mArco", "marco"))
+
+	require.False(t, hasPrefixCaseInsensitive("marco", "a"))
+	require.False(t, hasPrefixCaseInsensitive("marco", "abcdefghi"))
 }
 
 func TestHasSuffixCaseInsensitive(t *testing.T) {
-	assert.True(t, hasSuffixCaseInsensitive("marco", "rco"))
-	assert.True(t, hasSuffixCaseInsensitive("marco", "RcO"))
-	assert.True(t, hasSuffixCaseInsensitive("marco", "marco"))
-	assert.False(t, hasSuffixCaseInsensitive("marco", "a"))
-	assert.False(t, hasSuffixCaseInsensitive("marco", "abcdefghi"))
+	require.True(t, hasSuffixCaseInsensitive("marco", "rco"))
+	require.True(t, hasSuffixCaseInsensitive("marco", "RcO"))
+	require.True(t, hasSuffixCaseInsensitive("marco", "marco"))
+	require.False(t, hasSuffixCaseInsensitive("marco", "a"))
+	require.False(t, hasSuffixCaseInsensitive("marco", "abcdefghi"))
 }
 
 func getTestNameFromRegexp(re string) string {

diff --git a/promql/engine_test.go b/promql/engine_test.go
@@ -205,7 +205,7 @@ func (*errQuerier) LabelValues(context.Context, string, ...*labels.Matcher) ([]s
 	return nil, nil, nil
 }
 
-func (*errQuerier) LabelValuesStream(string, ...*labels.Matcher) storage.LabelValues {
+func (*errQuerier) LabelValuesStream(context.Context, string, ...*labels.Matcher) storage.LabelValues {
 	return nil
 }
 

diff --git a/storage/fanout_test.go b/storage/fanout_test.go
@@ -240,7 +240,7 @@ func (errQuerier) LabelValues(context.Context, string, ...*labels.Matcher) ([]st
 	return nil, nil, errors.New("label values error")
 }
 
-func (errQuerier) LabelValuesStream(string, ...*labels.Matcher) storage.LabelValues {
+func (errQuerier) LabelValuesStream(context.Context, string, ...*labels.Matcher) storage.LabelValues {
 	return storage.ErrLabelValues(errors.New("label values stream error"))
 }
 

diff --git a/storage/interface.go b/storage/interface.go
@@ -123,7 +123,7 @@ func (q *MockQuerier) LabelValues(context.Context, string, ...*labels.Matcher) (
 	return nil, nil, nil
 }
 
-func (q *MockQuerier) LabelValuesStream(string, ...*labels.Matcher) LabelValues {
+func (q *MockQuerier) LabelValuesStream(context.Context, string, ...*labels.Matcher) LabelValues {
 	return nil
 }
 
@@ -168,7 +168,7 @@ type LabelQuerier interface {
 	// It is not safe to use the strings beyond the lifetime of the querier.
 	// If matchers are specified the returned result set is reduced
 	// to label values of metrics matching the matchers.
-	LabelValuesStream(name string, matchers ...*labels.Matcher) LabelValues
+	LabelValuesStream(ctx context.Context, name string, matchers ...*labels.Matcher) LabelValues
 
 	// LabelNames returns all the unique label names present in the block in sorted order.
 	// If matchers are specified the returned result set is reduced

diff --git a/storage/merge.go b/storage/merge.go
@@ -219,18 +219,18 @@ func mergeStrings(a, b []string) []string {
 }
 
 // LabelValuesStream implements LabelQuerier.
-func (q *mergeGenericQuerier) LabelValuesStream(name string, matchers ...*labels.Matcher) LabelValues {
+func (q *mergeGenericQuerier) LabelValuesStream(ctx context.Context, name string, matchers ...*labels.Matcher) LabelValues {
 	if len(q.queriers) == 0 {
 		return EmptyLabelValues()
 	}
 	if len(q.queriers) == 1 {
-		return q.queriers[0].LabelValuesStream(name, matchers...)
+		return q.queriers[0].LabelValuesStream(ctx, name, matchers...)
 	}
 
 	h := make(labelValuesHeap, 0, len(q.queriers))
 	var ws annotations.Annotations
 	for _, sq := range q.queriers {
-		it := sq.LabelValuesStream(name, matchers...)
+		it := sq.LabelValuesStream(ctx, name, matchers...)
 		switch {
 		case it.Next():
 			h = append(h, it)

diff --git a/storage/merge_test.go b/storage/merge_test.go
@@ -1299,7 +1299,7 @@ func (m *mockGenericQuerier) LabelValues(_ context.Context, name string, matcher
 	return m.resp, m.warnings, m.err
 }
 
-func (m *mockGenericQuerier) LabelValuesStream(name string, matchers ...*labels.Matcher) LabelValues {
+func (m *mockGenericQuerier) LabelValuesStream(_ context.Context, name string, matchers ...*labels.Matcher) LabelValues {
 	m.mtx.Lock()
 	m.labelNamesRequested = append(m.labelNamesRequested, labelNameRequest{
 		name:     name,
@@ -1571,7 +1571,7 @@ func TestMergeGenericQuerierWithSecondaries_ErrorHandling(t *testing.T) {
 				}
 			})
 			t.Run("LabelValuesStream", func(t *testing.T) {
-				it := q.LabelValuesStream("test2")
+				it := q.LabelValuesStream(ctx, "test2")
 				require.NotNil(t, it)
 				var res []string
 				for it.Next() {

diff --git a/storage/noop.go b/storage/noop.go
@@ -35,7 +35,7 @@ func (noopQuerier) LabelValues(context.Context, string, ...*labels.Matcher) ([]s
 	return nil, nil, nil
 }
 
-func (noopQuerier) LabelValuesStream(string, ...*labels.Matcher) LabelValues {
+func (noopQuerier) LabelValuesStream(context.Context, string, ...*labels.Matcher) LabelValues {
 	return nil
 }
 
@@ -62,7 +62,7 @@ func (noopChunkQuerier) LabelValues(context.Context, string, ...*labels.Matcher)
 	return nil, nil, nil
 }
 
-func (noopChunkQuerier) LabelValuesStream(string, ...*labels.Matcher) LabelValues {
+func (noopChunkQuerier) LabelValuesStream(context.Context, string, ...*labels.Matcher) LabelValues {
 	return nil
 }
 

diff --git a/storage/remote/read.go b/storage/remote/read.go
@@ -216,7 +216,7 @@ func (q *querier) LabelValues(context.Context, string, ...*labels.Matcher) ([]st
 }
 
 // LabelValuesStream implements storage.LabelQuerier and is a noop.
-func (q *querier) LabelValuesStream(string, ...*labels.Matcher) storage.LabelValues {
+func (q *querier) LabelValuesStream(context.Context, string, ...*labels.Matcher) storage.LabelValues {
 	// TODO: Implement: https://github.com/prometheus/prometheus/issues/3351
 	return storage.ErrLabelValues(errors.New("not implemented"))
 }

diff --git a/storage/secondary.go b/storage/secondary.go
@@ -57,9 +57,9 @@ func (s *secondaryQuerier) LabelValues(ctx context.Context, name string, matcher
 	return vals, w, nil
 }
 
-func (s *secondaryQuerier) LabelValuesStream(name string, matchers ...*labels.Matcher) LabelValues {
+func (s *secondaryQuerier) LabelValuesStream(ctx context.Context, name string, matchers ...*labels.Matcher) LabelValues {
 	return &secondaryLabelValues{
-		it: s.genericQuerier.LabelValuesStream(name, matchers...),
+		it: s.genericQuerier.LabelValuesStream(ctx, name, matchers...),
 	}
 }