Alerting: Use URLs in image annotations (#66804)

* use tokens or urls in image annotations

* improve tests, fix some comments

* fix empty tokens

* code review changes, check for url before checking for token (support old token formats)

pkg/services/ngalert/notifier/images.go

@@ -2,7 +2,7 @@ package notifier
 
 import (
 	"context"
-	"errors"
+	"strings"
 
 	"github.com/grafana/alerting/images"
 
@@ -20,21 +20,27 @@ func newImageStore(store store.ImageStore) images.ImageStore {
 	}
 }
 
-func (i imageStore) GetImage(ctx context.Context, token string) (*images.Image, error) {
-	image, err := i.store.GetImage(ctx, token)
+func (i imageStore) GetImage(ctx context.Context, uri string) (*images.Image, error) {
+	var (
+		image *models.Image
+		err   error
+	)
+
+	// Check whether the uri is a URL or a token to know how to query the DB.
+	if strings.HasPrefix(uri, "http") {
+		image, err = i.store.GetImageByURL(ctx, uri)
+	} else {
+		token := strings.TrimPrefix(uri, "token://")
+		image, err = i.store.GetImage(ctx, token)
+	}
 	if err != nil {
-		if errors.Is(err, models.ErrImageNotFound) {
-			err = images.ErrImageNotFound
+		return nil, err
 	}
-	}
-	var result *images.Image
-	if image != nil {
-		result = &images.Image{
+
+	return &images.Image{
 		Token:     image.Token,
 		Path:      image.Path,
 		URL:       image.URL,
 		CreatedAt: image.CreatedAt,
-		}
-	}
-	return result, err
+	}, nil
 }

pkg/services/ngalert/notifier/images_test.go

@@ -0,0 +1,47 @@
+package notifier
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/ngalert/store"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetImage(t *testing.T) {
+	fakeImageStore := store.NewFakeImageStore(t)
+	store := newImageStore(fakeImageStore)
+
+	t.Run("queries by token when it gets a token", func(tt *testing.T) {
+		img := models.Image{
+			Token: "test",
+			URL:   "http://localhost:1234",
+			Path:  "test.png",
+		}
+		err := fakeImageStore.SaveImage(context.Background(), &img)
+		require.NoError(tt, err)
+
+		savedImg, err := store.GetImage(context.Background(), "token://"+img.Token)
+		require.NoError(tt, err)
+		require.Equal(tt, savedImg.Token, img.Token)
+		require.Equal(tt, savedImg.URL, img.URL)
+		require.Equal(tt, savedImg.Path, img.Path)
+	})
+
+	t.Run("queries by URL when it gets a URL", func(tt *testing.T) {
+		img := models.Image{
+			Token: "test",
+			Path:  "test.png",
+			URL:   "https://test.com/test.png",
+		}
+		err := fakeImageStore.SaveImage(context.Background(), &img)
+		require.NoError(tt, err)
+
+		savedImg, err := store.GetImage(context.Background(), img.URL)
+		require.NoError(tt, err)
+		require.Equal(tt, savedImg.Token, img.Token)
+		require.Equal(tt, savedImg.URL, img.URL)
+		require.Equal(tt, savedImg.Path, img.Path)
+	})
+}

pkg/services/ngalert/notifier/testing.go

@@ -31,6 +31,10 @@ func (f *fakeConfigStore) GetImage(ctx context.Context, token string) (*models.I
 	return nil, models.ErrImageNotFound
 }
 
+func (f *fakeConfigStore) GetImageByURL(ctx context.Context, url string) (*models.Image, error) {
+	return nil, models.ErrImageNotFound
+}
+
 func (f *fakeConfigStore) GetImages(ctx context.Context, tokens []string) ([]models.Image, []string, error) {
 	return nil, nil, models.ErrImageNotFound
 }

pkg/services/ngalert/schedule/compat.go

@@ -50,7 +50,7 @@ func stateToPostableAlert(alertState *state.State, appURL *url.URL) *models.Post
 	}
 
 	if alertState.Image != nil {
-		nA[alertingModels.ImageTokenAnnotation] = alertState.Image.Token
+		nA[alertingModels.ImageTokenAnnotation] = generateImageURI(alertState.Image)
 	}
 
 	if alertState.StateReason != "" {
@@ -167,3 +167,13 @@ func FromAlertsStateToStoppedAlert(firingStates []state.StateTransition, appURL
 	}
 	return alerts
 }
+
+// generateImageURI returns a string that serves as an identifier for the image.
+// It first checks if there is an image URL available, and if not,
+// it prefixes the image token with `token://` and uses it as the URI.
+func generateImageURI(image *ngModels.Image) string {
+	if image.URL != "" {
+		return image.URL
+	}
+	return "token://" + image.Token
+}

pkg/services/ngalert/schedule/compat_test.go

@@ -130,7 +130,7 @@ func Test_stateToPostableAlert(t *testing.T) {
 					for k, v := range alertState.Annotations {
 						expected[k] = v
 					}
-					expected["__alertImageToken__"] = alertState.Image.Token
+					expected["__alertImageToken__"] = "token://" + alertState.Image.Token
 
 					require.Equal(t, expected, result.Annotations)
 				})

pkg/services/ngalert/state/manager.go

@@ -422,8 +422,6 @@ func translateInstanceState(state ngModels.InstanceStateType) eval.State {
 func (st *Manager) deleteStaleStatesFromCache(ctx context.Context, logger log.Logger, evaluatedAt time.Time, alertRule *ngModels.AlertRule) []StateTransition {
 	// If we are removing two or more stale series it makes sense to share the resolved image as the alert rule is the same.
 	// TODO: We will need to change this when we support images without screenshots as each series will have a different image
-	var resolvedImage *ngModels.Image
-
 	staleStates := st.cache.deleteRuleStates(alertRule.GetKey(), func(s *State) bool {
 		return stateIsStale(evaluatedAt, s.LastEvaluationTime, alertRule.IntervalSeconds)
 	})
@@ -441,8 +439,6 @@ func (st *Manager) deleteStaleStatesFromCache(ctx context.Context, logger log.Lo
 
 		if oldState == eval.Alerting {
 			s.Resolved = true
-			// If there is no resolved image for this rule then take one
-			if resolvedImage == nil {
 			image, err := takeImage(ctx, st.images, alertRule)
 			if err != nil {
 				logger.Warn("Failed to take an image",
@@ -450,11 +446,9 @@ func (st *Manager) deleteStaleStatesFromCache(ctx context.Context, logger log.Lo
 					"panel", alertRule.GetPanelID(),
 					"error", err)
 			} else if image != nil {
-					resolvedImage = image
+				s.Image = image
 			}
 		}
-			s.Image = resolvedImage
-		}
 
 		record := StateTransition{
 			State:               s,

pkg/services/ngalert/state/state_test.go

@@ -584,7 +584,7 @@ func TestShouldTakeImage(t *testing.T) {
 		name:          "should not take image for alerting state with image",
 		state:         eval.Alerting,
 		previousState: eval.Alerting,
-		previousImage: &ngmodels.Image{Path: "foo.png", URL: "https://example.com/foo.png"},
+		previousImage: &ngmodels.Image{URL: "https://example.com/foo.png"},
 	}}
 
 	for _, test := range tests {

pkg/services/ngalert/store/image.go

@@ -20,6 +20,10 @@ type ImageStore interface {
 	// if the image has expired or if an image with the token does not exist.
 	GetImage(ctx context.Context, token string) (*models.Image, error)
 
+	// GetImageByURL looks for a image by its URL. It returns ErrImageNotFound
+	// if the image has expired or if there is no image associated with the URL.
+	GetImageByURL(ctx context.Context, url string) (*models.Image, error)
+
 	// GetImages returns all images that match the tokens. If one or more images
 	// have expired or do not exist then it also returns the unmatched tokens
 	// and an ErrImageNotFound error.
@@ -54,6 +58,23 @@ func (st DBstore) GetImage(ctx context.Context, token string) (*models.Image, er
 	return &image, nil
 }
 
+func (st DBstore) GetImageByURL(ctx context.Context, url string) (*models.Image, error) {
+	var image models.Image
+	if err := st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
+		exists, err := sess.Where("url = ? AND expires_at > ?", url, TimeNow().UTC()).Limit(1).Get(&image)
+		if err != nil {
+			return fmt.Errorf("failed to get image: %w", err)
+		} else if !exists {
+			return models.ErrImageNotFound
+		} else {
+			return nil
+		}
+	}); err != nil {
+		return nil, err
+	}
+	return &image, nil
+}
+
 func (st DBstore) GetImages(ctx context.Context, tokens []string) ([]models.Image, []string, error) {
 	var images []models.Image
 	if err := st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {

pkg/services/ngalert/store/image_test.go

@@ -30,7 +30,7 @@ func TestIntegrationSaveAndGetImage(t *testing.T) {
 	// create an image with a path on disk
 	image1 := models.Image{Path: "example.png"}
 	require.NoError(t, dbstore.SaveImage(ctx, &image1))
-	require.NotEqual(t, "", image1.Token)
+	require.NotEqual(t, image1.Token, "")
 
 	// image should not have expired
 	assert.False(t, image1.HasExpired())
@@ -49,7 +49,12 @@ func TestIntegrationSaveAndGetImage(t *testing.T) {
 	// create an image with a URL
 	image2 := models.Image{URL: "https://example.com/example.png"}
 	require.NoError(t, dbstore.SaveImage(ctx, &image2))
-	require.NotEqual(t, "", image2.Token)
+	require.NotEqual(t, image2.Token, "")
+
+	// create another image with the same URL
+	image3 := models.Image{URL: "https://example.com/example.png"}
+	require.NoError(t, dbstore.SaveImage(ctx, &image3))
+	require.NotEqual(t, image3.Token, "")
 
 	// image should not have expired
 	assert.False(t, image2.HasExpired())
@@ -60,12 +65,24 @@ func TestIntegrationSaveAndGetImage(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, image2, *result2)
 
+	// querying by URL should yield the same result even though we have two images with the same URL
+	result2, err = dbstore.GetImageByURL(ctx, image2.URL)
+	require.NoError(t, err)
+	assert.Equal(t, image2, *result2)
+
 	// expired image should not be returned
 	image1.ExpiresAt = time.Now().Add(-time.Second)
 	require.NoError(t, dbstore.SaveImage(ctx, &image1))
 	result1, err = dbstore.GetImage(ctx, image1.Token)
 	assert.EqualError(t, err, "image not found")
 	assert.Nil(t, result1)
+
+	// Querying by URL should yield the same result.
+	image2.ExpiresAt = time.Now().Add(-time.Second)
+	require.NoError(t, dbstore.SaveImage(ctx, &image1))
+	result2, err = dbstore.GetImage(ctx, image2.URL)
+	assert.EqualError(t, err, "image not found")
+	assert.Nil(t, result2)
 }
 
 func TestIntegrationGetImages(t *testing.T) {

pkg/services/ngalert/store/testing.go

@@ -49,6 +49,18 @@ func (s *FakeImageStore) GetImage(_ context.Context, token string) (*models.Imag
 	return nil, models.ErrImageNotFound
 }
 
+func (s *FakeImageStore) GetImageByURL(_ context.Context, url string) (*models.Image, error) {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	for _, image := range s.images {
+		if image.URL == url {
+			return image, nil
+		}
+	}
+
+	return nil, models.ErrImageNotFound
+}
+
 func (s *FakeImageStore) GetImages(_ context.Context, tokens []string) ([]models.Image, []string, error) {
 	s.mtx.Lock()
 	defer s.mtx.Unlock()