Implementing error handling on labels apis

Signed-off-by: Alan Protasio <alanprot@gmail.com>
cortexproject · Jul 3, 2023 · 7fbec0a · 7fbec0a
1 parent 465da07
commit 7fbec0a
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 2 deletions.
diff --git a/pkg/storegateway/bucket_stores.go b/pkg/storegateway/bucket_stores.go
@@ -335,7 +335,19 @@ func (u *BucketStores) LabelNames(ctx context.Context, req *storepb.LabelNamesRe
 		return &storepb.LabelNamesResponse{}, nil
 	}
 
-	return store.LabelNames(ctx, req)
+	err := u.getStoreError(userID)
+
+	if err != nil && errors.Is(err, bucket.ErrCustomerManagedKeyError) {
+		return nil, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", err)
+	}
+
+	resp, err := store.LabelNames(ctx, req)
+
+	if err != nil && errors.Is(err, bucket.ErrCustomerManagedKeyError) {
+		return resp, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", err)
+	}
+
+	return resp, err
 }
 
 // LabelValues implements the Storegateway proto service.
@@ -353,7 +365,19 @@ func (u *BucketStores) LabelValues(ctx context.Context, req *storepb.LabelValues
 		return &storepb.LabelValuesResponse{}, nil
 	}
 
-	return store.LabelValues(ctx, req)
+	err := u.getStoreError(userID)
+
+	if err != nil && errors.Is(err, bucket.ErrCustomerManagedKeyError) {
+		return nil, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", err)
+	}
+
+	resp, err := store.LabelValues(ctx, req)
+
+	if err != nil && errors.Is(err, bucket.ErrCustomerManagedKeyError) {
+		return resp, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", err)
+	}
+
+	return resp, err
 }
 
 // scanUsers in the bucket and return the list of found users. If an error occurs while

diff --git a/pkg/storegateway/bucket_stores_test.go b/pkg/storegateway/bucket_stores_test.go
@@ -84,8 +84,16 @@ func TestBucketStores_CustomerKeyError(t *testing.T) {
 
 	_, _, err = querySeries(stores, "user-1", "anything", 0, 100)
 	require.Equal(t, err, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", bucket.ErrCustomerManagedKeyError))
+	_, err = queryLabelsNames(stores, "user-1", "anything")
+	require.Equal(t, err, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", bucket.ErrCustomerManagedKeyError))
+	_, err = queryLabelsValues(stores, "user-1", "anything")
+	require.Equal(t, err, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", bucket.ErrCustomerManagedKeyError))
 	_, _, err = querySeries(stores, "user-2", "anything", 0, 100)
 	require.NoError(t, err)
+	_, err = queryLabelsNames(stores, "user-1", "anything")
+	require.Equal(t, err, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", bucket.ErrCustomerManagedKeyError))
+	_, err = queryLabelsValues(stores, "user-1", "anything")
+	require.Equal(t, err, httpgrpc.Errorf(int(codes.ResourceExhausted), "store error: %s", bucket.ErrCustomerManagedKeyError))
 
 	// Cleaning the error
 	mBucket.GetFailures = map[string]error{}
@@ -96,6 +104,10 @@ func TestBucketStores_CustomerKeyError(t *testing.T) {
 	require.NoError(t, err)
 	_, _, err = querySeries(stores, "user-2", "anything", 0, 100)
 	require.NoError(t, err)
+	_, err = queryLabelsNames(stores, "user-1", "anything")
+	require.NoError(t, err)
+	_, err = queryLabelsValues(stores, "user-1", "anything")
+	require.NoError(t, err)
 }
 
 func TestBucketStores_InitialSync(t *testing.T) {
@@ -487,6 +499,34 @@ func querySeries(stores *BucketStores, userID, metricName string, minT, maxT int
 	return srv.SeriesSet, srv.Warnings, err
 }
 
+func queryLabelsNames(stores *BucketStores, userID, metricName string) (*storepb.LabelNamesResponse, error) {
+	req := &storepb.LabelNamesRequest{
+		Matchers: []storepb.LabelMatcher{{
+			Type:  storepb.LabelMatcher_EQ,
+			Name:  labels.MetricName,
+			Value: metricName,
+		}},
+		PartialResponseStrategy: storepb.PartialResponseStrategy_ABORT,
+	}
+
+	ctx := setUserIDToGRPCContext(context.Background(), userID)
+	return stores.LabelNames(ctx, req)
+}
+
+func queryLabelsValues(stores *BucketStores, userID, metricName string) (*storepb.LabelValuesResponse, error) {
+	req := &storepb.LabelValuesRequest{
+		Matchers: []storepb.LabelMatcher{{
+			Type:  storepb.LabelMatcher_EQ,
+			Name:  labels.MetricName,
+			Value: metricName,
+		}},
+		PartialResponseStrategy: storepb.PartialResponseStrategy_ABORT,
+	}
+
+	ctx := setUserIDToGRPCContext(context.Background(), userID)
+	return stores.LabelValues(ctx, req)
+}
+
 func mockLoggingLevel() logging.Level {
 	level := logging.Level{}
 	err := level.Set("info")