Unified storage kvstore bulk import support (#113791)

* implement batchdelete in datastore

* implement bulkprocess in kv storage_backend

* convert bulkRVs to snowflake

pkg/storage/unified/resource/storage_backend.go

@@ -18,6 +18,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/trace"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 
 	"github.com/grafana/grafana/pkg/apimachinery/utils"
 	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
@@ -54,6 +55,7 @@ func convertEmptyToClusterNamespace(namespace string, withExperimentalClusterSco
 type kvStorageBackend struct {
 	snowflake                    *snowflake.Node
 	kv                           KV
+	bulkLock                     *BulkLock
 	dataStore                    *dataStore
 	eventStore                   *eventStore
 	notifier                     *notifier
@@ -102,6 +104,7 @@ func NewKVStorageBackend(opts KVBackendOptions) (StorageBackend, error) {
 
 	backend := &kvStorageBackend{
 		kv:                           kv,
+		bulkLock:                     NewBulkLock(),
 		dataStore:                    newDataStore(kv),
 		eventStore:                   eventStore,
 		notifier:                     newNotifier(eventStore, notifierOptions{}),
@@ -1236,6 +1239,187 @@ func (k *kvStorageBackend) GetResourceLastImportTimes(ctx context.Context) iter.
 	}
 }
 
+func (b *kvStorageBackend) ProcessBulk(ctx context.Context, setting BulkSettings, iter BulkRequestIterator) *resourcepb.BulkResponse {
+	// TODO cross-node lock
+	err := b.bulkLock.Start(setting.Collection)
+	if err != nil {
+		return &resourcepb.BulkResponse{
+			Error: AsErrorResult(err),
+		}
+	}
+	defer b.bulkLock.Finish(setting.Collection)
+
+	bulkRvGenerator := newBulkRV()
+	summaries := make(map[string]*resourcepb.BulkResponse_Summary, len(setting.Collection))
+	rsp := &resourcepb.BulkResponse{}
+
+	if setting.RebuildCollection {
+		for _, key := range setting.Collection {
+			events := make([]string, 0)
+			for evtKeyStr, err := range b.eventStore.ListKeysSince(ctx, 1) {
+				if err != nil {
+					b.log.Error("failed to list event: %s", err)
+					return rsp
+				}
+
+				evtKey, err := ParseEventKey(evtKeyStr)
+				if err != nil {
+					b.log.Error("error parsing event key: %s", err)
+					return rsp
+				}
+
+				if evtKey.Group != key.Group || evtKey.Resource != key.Resource || evtKey.Namespace != key.Namespace {
+					continue
+				}
+
+				events = append(events, evtKeyStr)
+			}
+
+			if err := b.eventStore.batchDelete(ctx, events); err != nil {
+				b.log.Error("failed to delete events: %s", err)
+				return rsp
+			}
+
+			historyKeys := make([]DataKey, 0)
+
+			for dataKey, err := range b.dataStore.Keys(ctx, ListRequestKey{
+				Namespace: key.Namespace,
+				Group:     key.Group,
+				Resource:  key.Resource,
+			}, SortOrderAsc) {
+				if err != nil {
+					b.log.Error("failed to list collection before delete: %s", err)
+					return rsp
+				}
+
+				historyKeys = append(historyKeys, dataKey)
+			}
+
+			previousCount := int64(len(historyKeys))
+			if err := b.dataStore.batchDelete(ctx, historyKeys); err != nil {
+				b.log.Error("failed to delete collection: %s", err)
+				return rsp
+			}
+			summaries[NSGR(key)] = &resourcepb.BulkResponse_Summary{
+				Namespace:     key.Namespace,
+				Group:         key.Group,
+				Resource:      key.Resource,
+				PreviousCount: previousCount,
+			}
+		}
+	} else {
+		for _, key := range setting.Collection {
+			summaries[NSGR(key)] = &resourcepb.BulkResponse_Summary{
+				Namespace: key.Namespace,
+				Group:     key.Group,
+				Resource:  key.Resource,
+			}
+		}
+	}
+
+	obj := &unstructured.Unstructured{}
+
+	saved := make([]DataKey, 0)
+	rollback := func() {
+		// we don't have transactions in the kv store, so we simply delete everything we created
+		err = b.dataStore.batchDelete(ctx, saved)
+		if err != nil {
+			b.log.Error("failed to delete during rollback: %s", err)
+		}
+	}
+
+	for iter.Next() {
+		if iter.RollbackRequested() {
+			rollback()
+			break
+		}
+
+		req := iter.Request()
+		if req == nil {
+			rollback()
+			rsp.Error = AsErrorResult(fmt.Errorf("missing request"))
+			break
+		}
+
+		rsp.Processed++
+
+		var action DataAction
+		switch resourcepb.WatchEvent_Type(req.Action) {
+		case resourcepb.WatchEvent_ADDED:
+			action = DataActionCreated
+			// Check if resource already exists for create operations
+			_, err := b.dataStore.GetLatestResourceKey(ctx, GetRequestKey{
+				Group:     req.Key.Group,
+				Resource:  req.Key.Resource,
+				Namespace: req.Key.Namespace,
+				Name:      req.Key.Name,
+			})
+			if err == nil {
+				rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{
+					Key:    req.Key,
+					Action: req.Action,
+					Error:  "resource already exists",
+				})
+				continue
+			}
+			if !errors.Is(err, ErrNotFound) {
+				rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{
+					Key:    req.Key,
+					Action: req.Action,
+					Error:  fmt.Sprintf("failed to check if resource exists: %s", err),
+				})
+				continue
+			}
+		case resourcepb.WatchEvent_MODIFIED:
+			action = DataActionUpdated
+		case resourcepb.WatchEvent_DELETED:
+			action = DataActionDeleted
+		default:
+			rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{
+				Key:    req.Key,
+				Action: req.Action,
+				Error:  "invalid event type",
+			})
+			continue
+		}
+
+		err := obj.UnmarshalJSON(req.Value)
+		if err != nil {
+			rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{
+				Key:    req.Key,
+				Action: req.Action,
+				Error:  "unable to unmarshal json",
+			})
+			continue
+		}
+
+		dataKey := DataKey{
+			Group:           req.Key.Group,
+			Resource:        req.Key.Resource,
+			Namespace:       req.Key.Namespace,
+			Name:            req.Key.Name,
+			ResourceVersion: bulkRvGenerator.next(obj),
+			Action:          action,
+			Folder:          req.Folder,
+		}
+		err = b.dataStore.Save(ctx, dataKey, bytes.NewReader(req.Value))
+		if err != nil {
+			rsp.Rejected = append(rsp.Rejected, &resourcepb.BulkResponse_Rejected{
+				Key:    req.Key,
+				Action: req.Action,
+				Error:  fmt.Sprintf("failed to save resource: %s", err),
+			})
+			continue
+		}
+
+		saved = append(saved, dataKey)
+	}
+
+	// TODO update last import time
+
+	return rsp
+}
+
 // readAndClose reads all data from a ReadCloser and ensures it's closed,
 // combining any errors from both operations.
 func readAndClose(r io.ReadCloser) ([]byte, error) {