From b0881daf234a0211cc5c486b2dd6402adf545695 Mon Sep 17 00:00:00 2001 From: Santiago Date: Wed, 26 Apr 2023 13:06:18 -0300 Subject: [PATCH] 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 | 38 +++++++++------- pkg/services/ngalert/notifier/images_test.go | 47 ++++++++++++++++++++ pkg/services/ngalert/notifier/testing.go | 4 ++ pkg/services/ngalert/schedule/compat.go | 12 ++++- pkg/services/ngalert/schedule/compat_test.go | 2 +- pkg/services/ngalert/state/manager.go | 22 ++++----- pkg/services/ngalert/state/state_test.go | 2 +- pkg/services/ngalert/store/image.go | 21 +++++++++ pkg/services/ngalert/store/image_test.go | 21 ++++++++- pkg/services/ngalert/store/testing.go | 12 +++++ 10 files changed, 146 insertions(+), 35 deletions(-) create mode 100644 pkg/services/ngalert/notifier/images_test.go diff --git a/pkg/services/ngalert/notifier/images.go b/pkg/services/ngalert/notifier/images.go index bbeec8b36bb..10a29b00663 100644 --- a/pkg/services/ngalert/notifier/images.go +++ b/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{ - Token: image.Token, - Path: image.Path, - URL: image.URL, - CreatedAt: image.CreatedAt, - } - } - return result, err + + return &images.Image{ + Token: image.Token, + Path: image.Path, + URL: image.URL, + CreatedAt: image.CreatedAt, + }, nil } diff --git a/pkg/services/ngalert/notifier/images_test.go b/pkg/services/ngalert/notifier/images_test.go new file mode 100644 index 00000000000..ba67ba52cb2 --- /dev/null +++ b/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) + }) +} diff --git a/pkg/services/ngalert/notifier/testing.go b/pkg/services/ngalert/notifier/testing.go index 8dd7e111765..31606030bf6 100644 --- a/pkg/services/ngalert/notifier/testing.go +++ b/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 } diff --git a/pkg/services/ngalert/schedule/compat.go b/pkg/services/ngalert/schedule/compat.go index 94f8a14b2a3..14c87ae663a 100644 --- a/pkg/services/ngalert/schedule/compat.go +++ b/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 +} diff --git a/pkg/services/ngalert/schedule/compat_test.go b/pkg/services/ngalert/schedule/compat_test.go index 419f79795ea..7e61fd3d055 100644 --- a/pkg/services/ngalert/schedule/compat_test.go +++ b/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) }) diff --git a/pkg/services/ngalert/state/manager.go b/pkg/services/ngalert/state/manager.go index 5161ba957cf..15e4ff66a52 100644 --- a/pkg/services/ngalert/state/manager.go +++ b/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,19 +439,15 @@ 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", - "dashboard", alertRule.GetDashboardUID(), - "panel", alertRule.GetPanelID(), - "error", err) - } else if image != nil { - resolvedImage = image - } + image, err := takeImage(ctx, st.images, alertRule) + if err != nil { + logger.Warn("Failed to take an image", + "dashboard", alertRule.GetDashboardUID(), + "panel", alertRule.GetPanelID(), + "error", err) + } else if image != nil { + s.Image = image } - s.Image = resolvedImage } record := StateTransition{ diff --git a/pkg/services/ngalert/state/state_test.go b/pkg/services/ngalert/state/state_test.go index 95ced7f8fe3..9d4684e45bc 100644 --- a/pkg/services/ngalert/state/state_test.go +++ b/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 { diff --git a/pkg/services/ngalert/store/image.go b/pkg/services/ngalert/store/image.go index 7d5a9f44b7b..095098138bb 100644 --- a/pkg/services/ngalert/store/image.go +++ b/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 { diff --git a/pkg/services/ngalert/store/image_test.go b/pkg/services/ngalert/store/image_test.go index 8002898d0ba..1f69e2ee7d1 100644 --- a/pkg/services/ngalert/store/image_test.go +++ b/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) { diff --git a/pkg/services/ngalert/store/testing.go b/pkg/services/ngalert/store/testing.go index 5480f443867..6d28dd7005e 100644 --- a/pkg/services/ngalert/store/testing.go +++ b/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()