Unified storage kvstore bulk import support (#113791)

* implement batchdelete in datastore * implement bulkprocess in kv storage_backend * convert bulkRVs to snowflake
2025-12-04 09:26:50 -05:00
parent 64a3c298cf
commit ae4d2324d6
4 changed files with 326 additions and 0 deletions
@@ -6,10 +6,14 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"sync"
+	"time"

 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc/metadata"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

 	authlib "github.com/grafana/authlib/types"

@@ -328,3 +332,84 @@ func (b *batchRunner) RollbackRequested() bool {
 	}
 	return false
 }
+
+type bulkRV struct {
+	max     int64
+	counter int64
+}
+
+// Used when executing a bulk import so that we can generate snowflake RVs in the past
+func newBulkRV() *bulkRV {
+	t := snowflakeFromTime(time.Now())
+	return &bulkRV{
+		max:     t,
+		counter: 0,
+	}
+}
+
+func (x *bulkRV) next(obj metav1.Object) int64 {
+	ts := snowflakeFromTime(obj.GetCreationTimestamp().Time)
+	anno := obj.GetAnnotations()
+	if anno != nil {
+		v := anno[utils.AnnoKeyUpdatedTimestamp]
+		t, err := time.Parse(time.RFC3339, v)
+		if err == nil {
+			ts = snowflakeFromTime(t)
+		}
+	}
+	if ts > x.max || ts < 0 {
+		ts = x.max
+	}
+
+	x.counter++
+	return ts + x.counter
+}
+
+type BulkLock struct {
+	running map[string]bool
+	mu      sync.Mutex
+}
+
+func NewBulkLock() *BulkLock {
+	return &BulkLock{
+		running: make(map[string]bool),
+	}
+}
+
+func (x *BulkLock) Start(keys []*resourcepb.ResourceKey) error {
+	x.mu.Lock()
+	defer x.mu.Unlock()
+
+	// First verify that it is not already running
+	ids := make([]string, len(keys))
+	for i, k := range keys {
+		id := NSGR(k)
+		if x.running[id] {
+			return &apierrors.StatusError{ErrStatus: metav1.Status{
+				Code:    http.StatusPreconditionFailed,
+				Message: "bulk export is already running",
+			}}
+		}
+		ids[i] = id
+	}
+
+	// Then add the keys to the lock
+	for _, k := range ids {
+		x.running[k] = true
+	}
+	return nil
+}
+
+func (x *BulkLock) Finish(keys []*resourcepb.ResourceKey) {
+	x.mu.Lock()
+	defer x.mu.Unlock()
+	for _, k := range keys {
+		delete(x.running, NSGR(k))
+	}
+}
+
+func (x *BulkLock) Active() bool {
+	x.mu.Lock()
+	defer x.mu.Unlock()
+	return len(x.running) > 0
+}
@@ -537,6 +537,27 @@ func (d *dataStore) Delete(ctx context.Context, key DataKey) error {
 	return d.kv.Delete(ctx, dataSection, key.String())
 }

+func (n *dataStore) batchDelete(ctx context.Context, keys []DataKey) error {
+	for len(keys) > 0 {
+		batch := keys
+		if len(batch) > dataBatchSize {
+			batch = batch[:dataBatchSize]
+		}
+
+		keys = keys[len(batch):]
+		stringKeys := make([]string, len(batch))
+		for _, dataKey := range batch {
+			stringKeys = append(stringKeys, dataKey.String())
+		}
+
+		if err := n.kv.BatchDelete(ctx, dataSection, stringKeys); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // ParseKey parses a string key into a DataKey struct
 func ParseKey(key string) (DataKey, error) {
 	parts := strings.Split(key, "/")
@@ -2950,6 +2950,42 @@ func TestDataStore_getGroupResources(t *testing.T) {
 	}
 }

+func TestDataStore_BatchDelete(t *testing.T) {
+	ds := setupTestDataStore(t)
+	ctx := context.Background()
+
+	keys := make([]DataKey, 95)
+	for i := 0; i < 95; i++ {
+		rv := node.Generate().Int64()
+		keys[i] = DataKey{
+			Namespace:       "test-namespace",
+			Group:           "test-group",
+			Resource:        "test-resource",
+			Name:            fmt.Sprintf("test-name-%d", i),
+			ResourceVersion: rv,
+			Action:          DataActionCreated,
+			Folder:          "test-folder",
+		}
+		content := fmt.Sprintf("test-value-%d", i)
+		err := ds.Save(ctx, keys[i], bytes.NewReader([]byte(content)))
+		require.NoError(t, err)
+	}
+
+	err := ds.batchDelete(ctx, keys)
+	require.NoError(t, err)
+
+	// Verify all events were deleted
+	for i := 0; i < 95; i++ {
+		_, err := ds.Get(ctx, DataKey{
+			Namespace: "test-namespace",
+			Group:     "test-group",
+			Resource:  "test-resource",
+			Name:      fmt.Sprintf("test-name-%d", i),
+		})
+		require.Error(t, err, "Resource should have been deleted")
+	}
+}
+
 func TestDataStore_BatchGet(t *testing.T) {
 	ds := setupTestDataStore(t)
 	ctx := context.Background()
@@ -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) {