Storage: Add support for sortBy selector (#80680)

* add support for sortBy field selector

* use label selectors instead of field selectors

* set entity_labels on create & update

* make entity server integration tests work

* test fixes

* be more consistent with handling of empty body, meta or status

* workaround for database is locked errors during migration

* fix double import of sqlite3

* rename functions and tidy up

* refactor update

* disable integration tests until we can fix the database locking issue

pkg/services/store/entity/sqlstash/querybuilder.go

@@ -6,15 +6,33 @@ import (
 	"github.com/grafana/grafana/pkg/services/sqlstore/migrator"
 )
 
+type Direction int
+
+const (
+	Ascending Direction = iota
+	Descending
+)
+
+func (d Direction) String() string {
+	if d == Descending {
+		return "DESC"
+	}
+	return "ASC"
+}
+
 type selectQuery struct {
 	dialect  migrator.Dialect
 	fields   []string // SELECT xyz
 	from     string   // FROM object
+	offset   int64
 	limit    int64
 	oneExtra bool
 
 	where []string
 	args  []any
+
+	orderBy   []string
+	direction []Direction
 }
 
 func (q *selectQuery) addWhere(f string, val ...any) {
@@ -53,6 +71,11 @@ func (q *selectQuery) addWhereIn(f string, vals []string) {
 	}
 }
 
+func (q *selectQuery) addOrderBy(field string, direction Direction) {
+	q.orderBy = append(q.orderBy, field)
+	q.direction = append(q.direction, direction)
+}
+
 func (q *selectQuery) toQuery() (string, []any) {
 	args := q.args
 	sb := strings.Builder{}
@@ -77,17 +100,27 @@ func (q *selectQuery) toQuery() (string, []any) {
 		}
 	}
 
-	if q.limit > 0 || q.oneExtra {
-		limit := q.limit
-		if limit < 1 {
-			limit = 20
-			q.limit = limit
+	if len(q.orderBy) > 0 && len(q.direction) == len(q.orderBy) {
+		sb.WriteString(" ORDER BY ")
+		for i, f := range q.orderBy {
+			if i > 0 {
+				sb.WriteString(",")
+			}
+			sb.WriteString(q.dialect.Quote(f))
+			sb.WriteString(" ")
+			sb.WriteString(q.direction[i].String())
 		}
-		if q.oneExtra {
-			limit = limit + 1
-		}
-		sb.WriteString(" LIMIT ?")
-		args = append(args, limit)
 	}
+
+	limit := q.limit
+	if limit < 1 {
+		limit = 20
+		q.limit = limit
+	}
+	if q.oneExtra {
+		limit = limit + 1
+	}
+	sb.WriteString(q.dialect.LimitOffset(limit, q.offset))
+
 	return sb.String(), args
 }

pkg/services/store/entity/sqlstash/sql_storage_server.go

@@ -3,10 +3,12 @@ package sqlstash
 import (
 	"context"
 	"database/sql"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"math/rand"
+	"slices"
 	"strings"
 	"time"
 
@@ -89,6 +91,7 @@ func (s *sqlEntityServer) getReadFields(r *entity.ReadEntityRequest) []string {
 		"origin", "origin_key", "origin_ts",
 		"meta",
 		"title", "slug", "description", "labels", "fields",
+		"message",
 	}
 
 	if r.WithBody {
@@ -134,6 +137,7 @@ func (s *sqlEntityServer) rowToEntity(ctx context.Context, rows *sql.Rows, r *en
 		&raw.Origin.Source, &raw.Origin.Key, &raw.Origin.Time,
 		&raw.Meta,
 		&raw.Title, &raw.Slug, &raw.Description, &labels, &fields,
+		&raw.Message,
 	}
 	if r.WithBody {
 		args = append(args, &raw.Body)
@@ -147,10 +151,6 @@ func (s *sqlEntityServer) rowToEntity(ctx context.Context, rows *sql.Rows, r *en
 		return nil, err
 	}
 
-	if raw.Origin.Source == "" {
-		raw.Origin = nil
-	}
-
 	// unmarshal json labels
 	if labels != "" {
 		if err := json.Unmarshal([]byte(labels), &raw.Labels); err != nil {
@@ -158,6 +158,17 @@ func (s *sqlEntityServer) rowToEntity(ctx context.Context, rows *sql.Rows, r *en
 		}
 	}
 
+	// set empty body, meta or status to nil
+	if raw.Body != nil && len(raw.Body) == 0 {
+		raw.Body = nil
+	}
+	if raw.Meta != nil && len(raw.Meta) == 0 {
+		raw.Meta = nil
+	}
+	if raw.Status != nil && len(raw.Status) == 0 {
+		raw.Status = nil
+	}
+
 	return raw, nil
 }
 
@@ -278,6 +289,10 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 	}
 
 	createdAt := r.Entity.CreatedAt
+	if createdAt < 1000 {
+		createdAt = time.Now().UnixMilli()
+	}
+
 	createdBy := r.Entity.CreatedBy
 	if createdBy == "" {
 		modifier, err := appcontext.User(ctx)
@@ -289,6 +304,7 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 		}
 		createdBy = store.GetUserIDString(modifier)
 	}
+
 	updatedAt := r.Entity.UpdatedAt
 	updatedBy := r.Entity.UpdatedBy
 
@@ -315,6 +331,10 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 		// generate guid for new entity
 		current.Guid = uuid.New().String()
 
+		// set created at/by
+		current.CreatedAt = createdAt
+		current.CreatedBy = createdBy
+
 		// parse provided key
 		key, err := entity.ParseKey(r.Entity.Key)
 		if err != nil {
@@ -350,6 +370,7 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 
 		etag := createContentsHash(current.Body, current.Meta, current.Status)
 		current.ETag = etag
+
 		current.UpdatedAt = updatedAt
 		current.UpdatedBy = updatedBy
 
@@ -365,18 +386,21 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 			s.log.Error("error marshalling labels", "msg", err.Error())
 			return err
 		}
+		current.Labels = r.Entity.Labels
 
 		fields, err := json.Marshal(r.Entity.Fields)
 		if err != nil {
 			s.log.Error("error marshalling fields", "msg", err.Error())
 			return err
 		}
+		current.Fields = r.Entity.Fields
 
 		errors, err := json.Marshal(r.Entity.Errors)
 		if err != nil {
 			s.log.Error("error marshalling errors", "msg", err.Error())
 			return err
 		}
+		current.Errors = r.Entity.Errors
 
 		if current.Origin == nil {
 			current.Origin = &entity.EntityOriginInfo{}
@@ -409,13 +433,13 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 			"group":            current.Group,
 			"resource":         current.Resource,
 			"name":             current.Name,
-			"created_at":       createdAt,
-			"created_by":       createdBy,
+			"created_at":       current.CreatedAt,
+			"created_by":       current.CreatedBy,
 			"group_version":    current.GroupVersion,
 			"folder":           current.Folder,
 			"slug":             current.Slug,
-			"updated_at":       updatedAt,
-			"updated_by":       updatedBy,
+			"updated_at":       current.UpdatedAt,
+			"updated_by":       current.UpdatedBy,
 			"body":             current.Body,
 			"meta":             current.Meta,
 			"status":           current.Status,
@@ -459,7 +483,7 @@ func (s *sqlEntityServer) Create(ctx context.Context, r *entity.CreateEntityRequ
 
 		rsp.Entity = current
 
-		return nil // s.writeSearchInfo(ctx, tx, current)
+		return s.setLabels(ctx, tx, current.Guid, current.Labels)
 	})
 	if err != nil {
 		s.log.Error("error creating entity", "msg", err.Error())
@@ -475,8 +499,11 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 		return nil, err
 	}
 
-	timestamp := time.Now().UnixMilli()
 	updatedAt := r.Entity.UpdatedAt
+	if updatedAt < 1000 {
+		updatedAt = time.Now().UnixMilli()
+	}
+
 	updatedBy := r.Entity.UpdatedBy
 	if updatedBy == "" {
 		modifier, err := appcontext.User(ctx)
@@ -488,9 +515,6 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 		}
 		updatedBy = store.GetUserIDString(modifier)
 	}
-	if updatedAt < 1000 {
-		updatedAt = timestamp
-	}
 
 	rsp := &entity.UpdateEntityResponse{
 		Entity: &entity.Entity{},
@@ -519,10 +543,7 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 
 		rsp.Entity.Guid = current.Guid
 
-		// Clear the labels+refs
-		if _, err := tx.Exec(ctx, "DELETE FROM entity_labels WHERE guid=?", rsp.Entity.Guid); err != nil {
-			return err
-		}
+		// Clear the refs
 		if _, err := tx.Exec(ctx, "DELETE FROM entity_ref WHERE guid=?", rsp.Entity.Guid); err != nil {
 			return err
 		}
@@ -553,6 +574,7 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 
 		etag := createContentsHash(current.Body, current.Meta, current.Status)
 		current.ETag = etag
+
 		current.UpdatedAt = updatedAt
 		current.UpdatedBy = updatedBy
 
@@ -568,18 +590,21 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 			s.log.Error("error marshalling labels", "msg", err.Error())
 			return err
 		}
+		current.Labels = r.Entity.Labels
 
 		fields, err := json.Marshal(r.Entity.Fields)
 		if err != nil {
 			s.log.Error("error marshalling fields", "msg", err.Error())
 			return err
 		}
+		current.Fields = r.Entity.Fields
 
 		errors, err := json.Marshal(r.Entity.Errors)
 		if err != nil {
 			s.log.Error("error marshalling errors", "msg", err.Error())
 			return err
 		}
+		current.Errors = r.Entity.Errors
 
 		if current.Origin == nil {
 			current.Origin = &entity.EntityOriginInfo{}
@@ -619,8 +644,8 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 			"group_version":    current.GroupVersion,
 			"folder":           current.Folder,
 			"slug":             current.Slug,
-			"updated_at":       updatedAt,
-			"updated_by":       updatedBy,
+			"updated_at":       current.UpdatedAt,
+			"updated_by":       current.UpdatedBy,
 			"body":             current.Body,
 			"meta":             current.Meta,
 			"status":           current.Status,
@@ -684,7 +709,7 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 
 		rsp.Entity = current
 
-		return nil // s.writeSearchInfo(ctx, tx, current)
+		return s.setLabels(ctx, tx, current.Guid, current.Labels)
 	})
 	if err != nil {
 		s.log.Error("error updating entity", "msg", err.Error())
@@ -694,23 +719,22 @@ func (s *sqlEntityServer) Update(ctx context.Context, r *entity.UpdateEntityRequ
 	return rsp, err
 }
 
-/*
-func (s *sqlEntityServer) writeSearchInfo(
-	ctx context.Context,
-	tx *session.SessionTx,
-	current *entity.Entity,
-) error {
-	// parent_key := current.getParentKey()
+func (s *sqlEntityServer) setLabels(ctx context.Context, tx *session.SessionTx, guid string, labels map[string]string) error {
+	s.log.Debug("setLabels", "guid", guid, "labels", labels)
 
-	// Add the labels rows
-	for k, v := range current.Labels {
+	// Clear the old labels
+	if _, err := tx.Exec(ctx, "DELETE FROM entity_labels WHERE guid=?", guid); err != nil {
+		return err
+	}
+
+	// Add the new labels
+	for k, v := range labels {
 		query, args, err := s.dialect.InsertQuery(
 			"entity_labels",
 			map[string]any{
-				"key":   current.Key,
+				"guid":  guid,
 				"label": k,
 				"value": v,
-				// "parent_key": parent_key,
 			},
 		)
 		if err != nil {
@@ -725,7 +749,6 @@ func (s *sqlEntityServer) writeSearchInfo(
 
 	return nil
 }
-*/
 
 func (s *sqlEntityServer) Delete(ctx context.Context, r *entity.DeleteEntityRequest) (*entity.DeleteEntityResponse, error) {
 	if err := s.Init(); err != nil {
@@ -816,7 +839,7 @@ func (s *sqlEntityServer) History(ctx context.Context, r *entity.EntityHistoryRe
 	rr := &entity.ReadEntityRequest{
 		Key:        r.Key,
 		WithBody:   true,
-		WithStatus: false,
+		WithStatus: true,
 	}
 
 	query, err := s.getReadSelect(rr)
@@ -879,6 +902,75 @@ func (s *sqlEntityServer) History(ctx context.Context, r *entity.EntityHistoryRe
 	return rsp, err
 }
 
+type ContinueToken struct {
+	Sort        []string `json:"s"`
+	StartOffset int64    `json:"o"`
+}
+
+func (c *ContinueToken) String() string {
+	b, _ := json.Marshal(c)
+	return base64.StdEncoding.EncodeToString(b)
+}
+
+func GetContinueToken(r *entity.EntityListRequest) (*ContinueToken, error) {
+	if r.NextPageToken == "" {
+		return nil, nil
+	}
+
+	continueVal, err := base64.StdEncoding.DecodeString(r.NextPageToken)
+	if err != nil {
+		return nil, fmt.Errorf("error decoding continue token")
+	}
+
+	t := &ContinueToken{}
+	err = json.Unmarshal(continueVal, t)
+	if err != nil {
+		return nil, err
+	}
+
+	if !slices.Equal(t.Sort, r.Sort) {
+		return nil, fmt.Errorf("sort order changed")
+	}
+
+	return t, nil
+}
+
+var sortByFields = []string{
+	"guid",
+	"key",
+	"namespace", "group", "group_version", "resource", "name", "folder",
+	"resource_version", "size", "etag",
+	"created_at", "created_by",
+	"updated_at", "updated_by",
+	"origin", "origin_key", "origin_ts",
+	"title", "slug", "description",
+}
+
+type SortBy struct {
+	Field     string
+	Direction Direction
+}
+
+func ParseSortBy(sort string) (*SortBy, error) {
+	sortBy := &SortBy{
+		Field:     "guid",
+		Direction: Ascending,
+	}
+
+	if strings.HasSuffix(sort, "_desc") {
+		sortBy.Field = sort[:len(sort)-5]
+		sortBy.Direction = Descending
+	} else {
+		sortBy.Field = sort
+	}
+
+	if !slices.Contains(sortByFields, sortBy.Field) {
+		return nil, fmt.Errorf("invalid sort field '%s', valid fields: %v", sortBy.Field, sortByFields)
+	}
+
+	return sortBy, nil
+}
+
 func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest) (*entity.EntityListResponse, error) {
 	if err := s.Init(); err != nil {
 		return nil, err
@@ -892,10 +984,6 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
 		return nil, fmt.Errorf("missing user in context")
 	}
 
-	if r.NextPageToken != "" || len(r.Sort) > 0 {
-		return nil, fmt.Errorf("not yet supported")
-	}
-
 	rr := &entity.ReadEntityRequest{
 		WithBody:   r.WithBody,
 		WithStatus: r.WithStatus,
@@ -909,6 +997,7 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
 		from:     "entity", // the table
 		args:     []any{},
 		limit:    r.Limit,
+		offset:   0,
 		oneExtra: true, // request one more than the limit (and show next token if it exists)
 	}
 
@@ -951,8 +1040,13 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
 		entityQuery.addWhere("folder", r.Folder)
 	}
 
-	if r.NextPageToken != "" {
-		entityQuery.addWhere("guid>?", r.NextPageToken)
+	// if we have a page token, use that to specify the first record
+	continueToken, err := GetContinueToken(r)
+	if err != nil {
+		return nil, err
+	}
+	if continueToken != nil {
+		entityQuery.offset = continueToken.StartOffset
 	}
 
 	if len(r.Labels) > 0 {
@@ -971,6 +1065,14 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
 
 		entityQuery.addWhereInSubquery("guid", query, args)
 	}
+	for _, sort := range r.Sort {
+		sortBy, err := ParseSortBy(sort)
+		if err != nil {
+			return nil, err
+		}
+		entityQuery.addOrderBy(sortBy.Field, sortBy.Direction)
+	}
+	entityQuery.addOrderBy("guid", Ascending)
 
 	query, args := entityQuery.toQuery()
 
@@ -990,8 +1092,11 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
 
 		// found more than requested
 		if int64(len(rsp.Results)) >= entityQuery.limit {
-			// TODO? this only works if we sort by guid
-			rsp.NextPageToken = result.Guid
+			continueToken := &ContinueToken{
+				Sort:        r.Sort,
+				StartOffset: entityQuery.offset + entityQuery.limit,
+			}
+			rsp.NextPageToken = continueToken.String()
 			break
 		}
 

pkg/services/store/entity/sqlstash/sql_storage_server_test.go

@@ -31,6 +31,7 @@ func TestCreate(t *testing.T) {
 				Name:      "set-minimum-uid",
 				Key:       "/playlist.grafana.app/playlists/default/set-minimum-uid",
 				CreatedBy: "set-minimum-creator",
+				Origin:    &entity.EntityOriginInfo{},
 			},
 			false,
 			true,
@@ -103,7 +104,7 @@ func TestCreate(t *testing.T) {
 			require.Equal(t, tc.ent.Status, read.Status)
 			require.Equal(t, tc.ent.Title, read.Title)
 			require.Equal(t, tc.ent.Size, read.Size)
-			require.Equal(t, tc.ent.CreatedAt, read.CreatedAt)
+			require.Greater(t, read.CreatedAt, int64(0))
 			require.Equal(t, tc.ent.CreatedBy, read.CreatedBy)
 			require.Equal(t, tc.ent.UpdatedAt, read.UpdatedAt)
 			require.Equal(t, tc.ent.UpdatedBy, read.UpdatedBy)