Avoid cache only if the warning message contains partial data error

Author: justinjung04

pkg/querier/partialdata/partia_data.go

@@ -8,14 +8,15 @@ import (
 type partialDataCtxKey struct{}
 
 var (
-	ctxKey              = &partialDataCtxKey{}
-	partialDataErrorMsg = "Query result may contain partial data."
+	ctxKey = &partialDataCtxKey{}
 )
 
+const ErrorMsg string = "Query result may contain partial data."
+
 type Error struct{}
 
 func (e Error) Error() string {
-	return partialDataErrorMsg
+	return ErrorMsg
 }
 
 func ContextWithPartialData(ctx context.Context, isEnabled bool) context.Context {

pkg/querier/tripperware/queryrange/results_cache.go

@@ -27,6 +27,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/chunk/cache"
 	"github.com/cortexproject/cortex/pkg/cortexpb"
 	"github.com/cortexproject/cortex/pkg/querier"
+	"github.com/cortexproject/cortex/pkg/querier/partialdata"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
@@ -296,8 +297,11 @@ func (s resultsCache) shouldCacheResponse(ctx context.Context, req tripperware.R
 		return false
 	}
 	if res, ok := r.(*tripperware.PrometheusResponse); ok {
-		if len(res.Warnings) > 0 {
-			return false
+		partialDataErr := partialdata.Error{}
+		for _, warning := range res.Warnings {
+			if warning == partialDataErr.Error() {
+				return false
+			}
 		}
 	}
 

pkg/querier/tripperware/queryrange/results_cache_test.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/chunk/cache"
 	"github.com/cortexproject/cortex/pkg/cortexpb"
+	"github.com/cortexproject/cortex/pkg/querier/partialdata"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
@@ -555,7 +556,7 @@ func TestShouldCache(t *testing.T) {
 			expected: false,
 		},
 		{
-			name:    "contains warning",
+			name:    "contains partial data warning",
 			request: &tripperware.PrometheusRequest{Query: "metric"},
 			input: tripperware.Response(&tripperware.PrometheusResponse{
 				Headers: []*tripperware.PrometheusResponseHeader{
@@ -564,10 +565,24 @@ func TestShouldCache(t *testing.T) {
 						Values: []string{},
 					},
 				},
-				Warnings: []string{"some warning"},
+				Warnings: []string{partialdata.ErrorMsg},
 			}),
 			expected: false,
 		},
+		{
+			name:    "contains other warning",
+			request: &tripperware.PrometheusRequest{Query: "metric"},
+			input: tripperware.Response(&tripperware.PrometheusResponse{
+				Headers: []*tripperware.PrometheusResponseHeader{
+					{
+						Name:   "meaninglessheader",
+						Values: []string{},
+					},
+				},
+				Warnings: []string{"other warning"},
+			}),
+			expected: true,
+		},
 	} {
 		{
 			t.Run(tc.name, func(t *testing.T) {