From 5c19011dc14a2207a7b2b92db54e285a3ac2099a Mon Sep 17 00:00:00 2001
From: "Grot (@grafanabot)" <43478413+grafanabot@users.noreply.github.com>
Date: Tue, 31 May 2022 23:10:03 -0400
Subject: [PATCH] Alerting: Add GetImages to ImageStore (#49717) (#49791)

GetImages does a `TOKEN IN` query for each token in the argument.

(cherry picked from commit 47a3ddd968eda8d2cd09edd306cbb78ab8966590)

Co-authored-by: George Robinson <george.robinson@grafana.com>
---
 pkg/services/ngalert/notifier/testing.go |  4 ++
 pkg/services/ngalert/store/image.go      | 21 +++++++++-
 pkg/services/ngalert/store/image_test.go | 50 ++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/pkg/services/ngalert/notifier/testing.go b/pkg/services/ngalert/notifier/testing.go
index a1fc49a03e5..a7cd451b11c 100644
--- a/pkg/services/ngalert/notifier/testing.go
+++ b/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()
 
diff --git a/pkg/services/ngalert/store/image.go b/pkg/services/ngalert/store/image.go
index 4d0f94b0ee0..3f707cd463e 100644
--- a/pkg/services/ngalert/store/image.go
+++ b/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
diff --git a/pkg/services/ngalert/store/image_test.go b/pkg/services/ngalert/store/image_test.go
index bd607a9137f..52d6772c6c8 100644
--- a/pkg/services/ngalert/store/image_test.go
+++ b/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)