Alerting: Add GetImages to ImageStore (#49717)

* Alerting: Add GetImages to ImageStore

* Use assert.ElementsMatch instead of sort.Sort

pkg/services/ngalert/notifier/testing.go

@@ -27,6 +27,10 @@ func (f *FakeConfigStore) GetImage(ctx context.Context, token string) (*models.I
 	return nil, models.ErrImageNotFound
 }
 
+func (f *FakeConfigStore) GetImages(ctx context.Context, tokens []string) ([]models.Image, error) {
+	return nil, models.ErrImageNotFound
+}
+
 func NewFakeConfigStore(t *testing.T, configs map[int64]*models.AlertConfiguration) FakeConfigStore {
 	t.Helper()
 

pkg/services/ngalert/store/image.go

@@ -12,10 +12,14 @@ import (
 )
 
 type ImageStore interface {
-	// Get returns the image with the token or ErrImageNotFound.
+	// GetImage returns the image with the token or ErrImageNotFound.
 	GetImage(ctx context.Context, token string) (*models.Image, error)
 
-	// Saves the image or returns an error.
+	// GetImages returns all images that match the tokens. If one or more
+	// tokens does not exist then it also returns ErrImageNotFound.
+	GetImages(ctx context.Context, tokens []string) ([]models.Image, error)
+
+	// SaveImage saves the image or returns an error.
 	SaveImage(ctx context.Context, img *models.Image) error
 }
 
@@ -36,6 +40,19 @@ func (st DBstore) GetImage(ctx context.Context, token string) (*models.Image, er
 	return &img, nil
 }
 
+func (st DBstore) GetImages(ctx context.Context, tokens []string) ([]models.Image, error) {
+	var imgs []models.Image
+	if err := st.SQLStore.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		return sess.In("token", tokens).Find(&imgs)
+	}); err != nil {
+		return nil, err
+	}
+	if len(imgs) < len(tokens) {
+		return imgs, models.ErrImageNotFound
+	}
+	return imgs, nil
+}
+
 func (st DBstore) SaveImage(ctx context.Context, img *models.Image) error {
 	return st.SQLStore.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		// TODO: Is this a good idea? Do we actually want to automatically expire

pkg/services/ngalert/store/image_test.go

@@ -93,6 +93,56 @@ func TestIntegrationSaveAndGetImage(t *testing.T) {
 	}
 }
 
+func TestIntegrationGetImages(t *testing.T) {
+	mockTimeNow()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_, dbstore := tests.SetupTestEnv(t, baseIntervalSeconds)
+
+	// create an image foo.png
+	img1 := models.Image{Path: "foo.png"}
+	require.NoError(t, dbstore.SaveImage(ctx, &img1))
+
+	// GetImages should return the first image
+	imgs, err := dbstore.GetImages(ctx, []string{img1.Token})
+	require.NoError(t, err)
+	assert.Equal(t, []models.Image{img1}, imgs)
+
+	// create another image bar.png
+	img2 := models.Image{Path: "bar.png"}
+	require.NoError(t, dbstore.SaveImage(ctx, &img2))
+
+	// GetImages should return both images
+	imgs, err = dbstore.GetImages(ctx, []string{img1.Token, img2.Token})
+	require.NoError(t, err)
+	assert.ElementsMatch(t, []models.Image{img1, img2}, imgs)
+
+	// GetImages should return the first image
+	imgs, err = dbstore.GetImages(ctx, []string{img1.Token})
+	require.NoError(t, err)
+	assert.Equal(t, []models.Image{img1}, imgs)
+
+	// GetImages should return the second image
+	imgs, err = dbstore.GetImages(ctx, []string{img2.Token})
+	require.NoError(t, err)
+	assert.Equal(t, []models.Image{img2}, imgs)
+
+	// GetImages should return the first image and an error
+	imgs, err = dbstore.GetImages(ctx, []string{img1.Token, "unknown"})
+	assert.EqualError(t, err, "image not found")
+	assert.Equal(t, []models.Image{img1}, imgs)
+
+	// GetImages should return no images for no tokens
+	imgs, err = dbstore.GetImages(ctx, []string{})
+	require.NoError(t, err)
+	assert.Len(t, imgs, 0)
+
+	// GetImages should return no images for nil tokens
+	imgs, err = dbstore.GetImages(ctx, nil)
+	require.NoError(t, err)
+	assert.Len(t, imgs, 0)
+}
+
 func TestIntegrationDeleteExpiredImages(t *testing.T) {
 	mockTimeNow()
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)