Create context inside lazy iterator

Author: aaron7

pkg/chunk/chunk_store.go

@@ -262,7 +262,11 @@ outer:
 			}
 		}
 
-		newIterator, err := NewLazySeriesIterator(c, metric, from, through)
+		orgID, err := user.ExtractOrgID(ctx)
+		if err != nil {
+			return nil, err
+		}
+		newIterator, err := NewLazySeriesIterator(c, metric, from, through, orgID)
 		if err != nil {
 			return nil, err
 		}

pkg/chunk/chunk_store_test.go

@@ -299,7 +299,7 @@ func TestChunkStore_Get_lazy(t *testing.T) {
 				// Create expected iterators with current schema store
 				var expectedIterators []local.SeriesIterator
 				for _, expectedMetric := range tc.expectedIteratorMetrics {
-					newIterator, err := NewLazySeriesIterator(store, expectedMetric, from, now)
+					newIterator, err := NewLazySeriesIterator(store, expectedMetric, from, now, userID)
 					require.NoError(t, err)
 					expectedIterators = append(expectedIterators, newIterator)
 				}

pkg/chunk/iterator.go

@@ -9,6 +9,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/storage/local"
 	"github.com/prometheus/prometheus/storage/metric"
+	"github.com/weaveworks/common/user"
 )
 
 // LazySeriesIterator is a struct and not just a renamed type because otherwise the Metric
@@ -27,6 +28,7 @@ type LazySeriesIterator struct {
 	// be fetched. Use sync.Once to ensure the iterator is only created once.
 	sampleSeriesIterator *local.SeriesIterator
 	onceCreateIterator   sync.Once
+	orgID                string
 }
 
 type byMatcherLabel metric.LabelMatchers
@@ -36,14 +38,18 @@ func (lms byMatcherLabel) Swap(i, j int)      { lms[i], lms[j] = lms[j], lms[i]
 func (lms byMatcherLabel) Less(i, j int) bool { return lms[i].Name < lms[j].Name }
 
 // NewLazySeriesIterator creates a LazySeriesIterator.
-func NewLazySeriesIterator(chunkStore *Store, seriesMetric model.Metric, from model.Time, through model.Time) (*LazySeriesIterator, error) {
+func NewLazySeriesIterator(chunkStore *Store, seriesMetric model.Metric, from model.Time, through model.Time, orgID string) (*LazySeriesIterator, error) {
 	_, ok := seriesMetric[model.MetricNameLabel]
 	if !ok {
 		return nil, fmt.Errorf("series does not have a metric name")
 	}
 
 	var matchers metric.LabelMatchers
 	for labelName, labelValue := range seriesMetric {
+		if labelName == "__name__" {
+			continue
+		}
+
 		matcher, err := metric.NewLabelMatcher(metric.Equal, labelName, labelValue)
 		if err != nil {
 			return nil, err
@@ -58,6 +64,7 @@ func NewLazySeriesIterator(chunkStore *Store, seriesMetric model.Metric, from mo
 		from:       from,
 		through:    through,
 		matchers:   matchers,
+		orgID:      orgID,
 	}, nil
 }
 
@@ -87,6 +94,7 @@ func (it *LazySeriesIterator) RangeValues(in metric.Interval) []model.SamplePair
 	})
 	if err != nil {
 		// TODO: Handle error.
+		fmt.Printf("ERROR %+v", err)
 		return nil
 	}
 	return (*it.sampleSeriesIterator).RangeValues(in)
@@ -101,7 +109,7 @@ func (it *LazySeriesIterator) createSampleSeriesIterator() error {
 		return fmt.Errorf("series does not have a metric name")
 	}
 
-	ctx := context.Background()
+	ctx := user.InjectOrgID(context.Background(), it.orgID)
 	sampleSeriesIterators, err := it.chunkStore.getMetricNameIterators(ctx, it.from, it.through, it.matchers, metricName)
 	if err != nil {
 		return err

pkg/chunk/iterator_test.go

@@ -1,18 +1,20 @@
 package chunk
 
 import (
+	"context"
 	"testing"
 
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/storage/metric"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/user"
 )
 
 func TestLazySeriesIterator_Metric(t *testing.T) {
 	store := newTestChunkStore(t, StoreConfig{})
 	now := model.Now()
 	sampleMetric := model.Metric{model.MetricNameLabel: "foo"}
-	iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now)
+	iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now, userID)
 	require.NoError(t, err)
 	for _, c := range []struct {
 		iterator       *LazySeriesIterator
@@ -29,45 +31,91 @@ func TestLazySeriesIterator_Metric(t *testing.T) {
 }
 
 func TestLazySeriesIterator_ValueAtOrBeforeTime(t *testing.T) {
-	store := newTestChunkStore(t, StoreConfig{})
 	now := model.Now()
+	ctx := user.InjectOrgID(context.Background(), userID)
+
 	sampleMetric := model.Metric{model.MetricNameLabel: "foo"}
-	iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now)
+	dummyChunk := dummyChunkFor(sampleMetric)
+	dummySamples, err := dummyChunk.Samples()
 	require.NoError(t, err)
-	for _, c := range []struct {
-		iterator       *LazySeriesIterator
-		timestamp      model.Time
-		expectedSample model.SamplePair
+
+	schemas := []struct {
+		name string
+		fn   func(cfg SchemaConfig) Schema
 	}{
-		{
-			iterator:       iterator,
-			timestamp:      now,
-			expectedSample: model.ZeroSamplePair,
-		},
-	} {
-		sample := c.iterator.ValueAtOrBeforeTime(c.timestamp)
-		require.Equal(t, c.expectedSample, sample)
+		{"v8 schema", v8Schema},
+	}
+
+	for _, schema := range schemas {
+		// Create store with the dummy chunk.
+		store := newTestChunkStore(t, StoreConfig{schemaFactory: schema.fn})
+		store.Put(ctx, []Chunk{dummyChunk})
+
+		// Create the lazy series iterator
+		iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now, userID)
+		require.NoError(t, err)
+		for _, tc := range []struct {
+			iterator       *LazySeriesIterator
+			timestamp      model.Time
+			expectedSample model.SamplePair
+		}{
+			{
+				iterator:       iterator,
+				timestamp:      now,
+				expectedSample: dummySamples[0],
+			},
+		} {
+			// sampleSeriesIterator should be created lazily only when RangeValues is called.
+			require.Nil(t, tc.iterator.sampleSeriesIterator)
+			sample := tc.iterator.ValueAtOrBeforeTime(tc.timestamp)
+			require.NotNil(t, tc.iterator.sampleSeriesIterator)
+
+			require.Equal(t, tc.expectedSample, sample)
+		}
 	}
 }
 
 func TestLazySeriesIterator_RangeValues(t *testing.T) {
-	store := newTestChunkStore(t, StoreConfig{})
 	now := model.Now()
+	ctx := user.InjectOrgID(context.Background(), userID)
+
 	sampleMetric := model.Metric{model.MetricNameLabel: "foo"}
-	iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now)
+	dummyChunk := dummyChunkFor(sampleMetric)
+	dummySamples, err := dummyChunk.Samples()
 	require.NoError(t, err)
-	for _, c := range []struct {
-		iterator        *LazySeriesIterator
-		interval        metric.Interval
-		expectedSamples []model.SamplePair
+
+	schemas := []struct {
+		name string
+		fn   func(cfg SchemaConfig) Schema
 	}{
-		{
-			iterator:        iterator,
-			interval:        metric.Interval{OldestInclusive: now, NewestInclusive: now},
-			expectedSamples: nil,
-		},
-	} {
-		samples := c.iterator.RangeValues(c.interval)
-		require.Equal(t, c.expectedSamples, samples)
+		{"v8 schema", v8Schema},
+	}
+
+	for _, schema := range schemas {
+		// Create store with the dummy chunk.
+		store := newTestChunkStore(t, StoreConfig{schemaFactory: schema.fn})
+		store.Put(ctx, []Chunk{dummyChunk})
+
+		// Create the lazy series iterator
+		iterator, err := NewLazySeriesIterator(store, sampleMetric, now, now, userID)
+		require.NoError(t, err)
+		for _, tc := range []struct {
+			iterator        *LazySeriesIterator
+			interval        metric.Interval
+			expectedSamples []model.SamplePair
+		}{
+			{
+				iterator:        iterator,
+				interval:        metric.Interval{OldestInclusive: now, NewestInclusive: now},
+				expectedSamples: dummySamples,
+			},
+		} {
+			// sampleSeriesIterator should be created lazily only when RangeValues is called.
+			require.Nil(t, tc.iterator.sampleSeriesIterator)
+			samples := tc.iterator.RangeValues(tc.interval)
+			require.NotNil(t, tc.iterator.sampleSeriesIterator)
+
+			require.Equal(t, tc.expectedSamples, samples)
+		}
 	}
 }