grafana/pkg/services/ngalert/store/image.go

package store

import (
	"context"
	"fmt"
	"time"

	"github.com/google/uuid"

	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/sqlstore"
)

const (
	imageExpirationDuration = 24 * time.Hour
)

type ImageStore interface {
	// GetImage returns the image with the token. It returns ErrImageNotFound
	// if the image has expired or if an image with the token does not exist.
	GetImage(ctx context.Context, token 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.
	GetImages(ctx context.Context, tokens []string) ([]models.Image, []string, error)

	// SaveImage saves the image or returns an error.
	SaveImage(ctx context.Context, img *models.Image) error
}

type ImageAdminStore interface {
	ImageStore

	// DeleteExpiredImages deletes expired images. It returns the number of deleted images
	// or an error.
	DeleteExpiredImages(context.Context) (int64, error)
}

func (st DBstore) GetImage(ctx context.Context, token string) (*models.Image, error) {
	var image models.Image
	if err := st.SQLStore.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
		exists, err := sess.Where("token = ? AND expires_at > ?", token, TimeNow().UTC()).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 *sqlstore.DBSession) error {
		return sess.In("token", tokens).Where("expires_at > ?", TimeNow().UTC()).Find(&images)
	}); err != nil {
		return nil, nil, err
	}
	if len(images) < len(tokens) {
		return images, unmatchedTokens(tokens, images), models.ErrImageNotFound
	}
	return images, nil, nil
}

func (st DBstore) SaveImage(ctx context.Context, img *models.Image) error {
	return st.SQLStore.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
		if img.ID == 0 {
			// If the ID is zero then this is a new image. It needs a token, a created timestamp
			// and an expiration time. The expiration time of the image is derived from the created
			// timestamp rather than the current time as it helps assert that the expiration time
			// has the intended duration in tests.
			token, err := uuid.NewRandom()
			if err != nil {
				return fmt.Errorf("failed to create token: %w", err)
			}
			img.Token = token.String()
			img.CreatedAt = TimeNow().UTC()
			img.ExpiresAt = img.CreatedAt.Add(imageExpirationDuration)
			if _, err := sess.Insert(img); err != nil {
				return fmt.Errorf("failed to insert image: %w", err)
			}
		} else {
			// Check if the image exists as some databases return 0 rows affected if
			// no changes were made
			if ok, err := sess.Where("id = ?", img.ID).ForUpdate().Exist(&models.Image{}); err != nil {
				return fmt.Errorf("failed to check if image exists: %v", err)
			} else if !ok {
				return models.ErrImageNotFound
			}

			// Do not reset the expiration time as it can be extended with ExtendDuration
			if _, err := sess.ID(img.ID).Update(img); err != nil {
				return fmt.Errorf("failed to update image: %v", err)
			}
		}
		return nil
	})
}

func (st DBstore) DeleteExpiredImages(ctx context.Context) (int64, error) {
	var n int64
	if err := st.SQLStore.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
		rows, err := sess.Where("expires_at < ?", TimeNow().UTC()).Delete(&models.Image{})
		if err != nil {
			return fmt.Errorf("failed to delete expired images: %w", err)
		}
		n = rows
		return nil
	}); err != nil {
		return -1, err
	}
	return n, nil
}

// unmatchedTokens returns the tokens that were not matched to an image.
func unmatchedTokens(tokens []string, images []models.Image) []string {
	matched := make(map[string]struct{})
	for _, image := range images {
		matched[image.Token] = struct{}{}
	}
	unmatched := make([]string, 0, len(tokens))
	for _, token := range tokens {
		if _, ok := matched[token]; !ok {
			unmatched = append(unmatched, token)
		}
	}
	return unmatched
}