Cloud migrations: store snapshots in the database (#108551)

* Cloud migrations: store snapshots in the database

* update github.com/grafana/grafana-cloud-migration-snapshot to v1.9.0

* make update-workspace

* use new field name in test

* return error after call to fmt.Errorf

* create methods for readability / fix session deletiong not deleting snapshots

* remove debugging changes

* update sample.ini

* update tests to include OrgID in ListSnapshotsQuery

* lint

* lint

* Update pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

Co-authored-by: Matheus Macabu <macabu@users.noreply.github.com>

* remove TODO

* Update pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

Co-authored-by: Matheus Macabu <macabu@users.noreply.github.com>

* remove one of the debug logs

---------

Co-authored-by: Matheus Macabu <macabu@users.noreply.github.com>

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go

@@ -496,23 +496,24 @@ func (s *Service) CreateSnapshot(ctx context.Context, signedInUser *user.SignedI
 
 	// save snapshot to the db
 	snapshot := cloudmigration.CloudMigrationSnapshot{
-		UID:            util.GenerateShortUID(),
-		SessionUID:     cmd.SessionUID,
-		Status:         cloudmigration.SnapshotStatusCreating,
-		EncryptionKey:  initResp.EncryptionKey,
-		GMSSnapshotUID: initResp.SnapshotID,
-		LocalDir:       filepath.Join(s.cfg.CloudMigration.SnapshotFolder, "grafana", "snapshots", initResp.SnapshotID),
+		UID:                 util.GenerateShortUID(),
+		SessionUID:          cmd.SessionUID,
+		Status:              cloudmigration.SnapshotStatusCreating,
+		GMSPublicKey:        initResp.GMSPublicKey,
+		GMSSnapshotUID:      initResp.SnapshotID,
+		Metadata:            initResp.Metadata,
+		EncryptionAlgo:      initResp.Algo,
+		LocalDir:            filepath.Join(s.cfg.CloudMigration.SnapshotFolder, "grafana", "snapshots", initResp.SnapshotID),
+		ResourceStorageType: s.cfg.CloudMigration.ResourceStorageType,
 	}
 
-	uid, err := s.store.CreateSnapshot(ctx, snapshot)
-	if err != nil {
+	if err := s.store.CreateSnapshot(ctx, snapshot); err != nil {
 		return nil, fmt.Errorf("saving snapshot: %w", err)
 	}
-	snapshot.UID = uid
 
 	// Update status to "creating" to ensure the frontend polls from now on
 	if err := s.updateSnapshotWithRetries(ctx, cloudmigration.UpdateSnapshotCmd{
-		UID:       uid,
+		UID:       snapshot.UID,
 		SessionID: cmd.SessionUID,
 		Status:    cloudmigration.SnapshotStatusCreating,
 	}); err != nil {
@@ -538,6 +539,7 @@ func (s *Service) CreateSnapshot(ctx context.Context, signedInUser *user.SignedI
 		s.report(asyncCtx, session, gmsclient.EventStartBuildingSnapshot, 0, nil, signedInUser.UserUID)
 
 		start := time.Now()
+
 		err := s.buildSnapshot(asyncCtx, signedInUser, initResp.MaxItemsPerPartition, initResp.Metadata, snapshot, cmd.ResourceTypes)
 		if err != nil {
 			asyncSpan.SetStatus(codes.Error, "error building snapshot")
@@ -896,10 +898,12 @@ func (s *Service) deleteLocalFiles(snapshots []cloudmigration.CloudMigrationSnap
 
 	var err error
 	for _, snapshot := range snapshots {
-		err = os.RemoveAll(snapshot.LocalDir)
-		if err != nil {
-			// in this case we only log the error, don't return it to continue with the process
-			s.log.Error("deleting migration snapshot files", "err", err)
+		if snapshot.LocalDir != "" {
+			err = os.RemoveAll(snapshot.LocalDir)
+			if err != nil {
+				// in this case we only log the error, don't return it to continue with the process
+				s.log.Error("deleting migration snapshot files", "err", err)
+			}
 		}
 	}
 	return err

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration_test.go

@@ -103,14 +103,15 @@ func Test_GetSnapshotStatusFromGMS(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		uid, err := s.store.CreateSnapshot(ctx, cloudmigration.CloudMigrationSnapshot{
-			UID:            "test uid",
+		uid := "test uid"
+
+		err = s.store.CreateSnapshot(ctx, cloudmigration.CloudMigrationSnapshot{
+			UID:            uid,
 			SessionUID:     sess.UID,
 			Status:         cloudmigration.SnapshotStatusCreating,
 			GMSSnapshotUID: "gms uid",
 		})
 		require.NoError(t, err)
-		assert.Equal(t, "test uid", uid)
 
 		// Make sure status is coming from the db only
 		snapshot, err := s.GetSnapshot(ctx, cloudmigration.GetSnapshotsQuery{
@@ -381,8 +382,9 @@ func Test_OnlyQueriesStatusFromGMSWhenRequired(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	uid, err := s.store.CreateSnapshot(context.Background(), cloudmigration.CloudMigrationSnapshot{
-		UID:            uuid.NewString(),
+	uid := uuid.NewString()
+	err = s.store.CreateSnapshot(context.Background(), cloudmigration.CloudMigrationSnapshot{
+		UID:            uid,
 		SessionUID:     sess.UID,
 		Status:         cloudmigration.SnapshotStatusCreating,
 		GMSSnapshotUID: "gms uid",

pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

@@ -1,6 +1,7 @@
 package cloudmigrationimpl
 
 import (
+	"bytes"
 	"context"
 	cryptoRand "crypto/rand"
 	"encoding/json"
@@ -562,7 +563,7 @@ func (s *Service) buildSnapshot(
 
 	// Use GMS public key + the grafana generated private key to encrypt snapshot files.
 	snapshotWriter, err := snapshot.NewSnapshotWriter(contracts.AssymetricKeys{
-		Public:  snapshotMeta.EncryptionKey,
+		Public:  snapshotMeta.GMSPublicKey,
 		Private: privateKey[:],
 	},
 		crypto.NewNacl(),
@@ -605,6 +606,56 @@ func (s *Service) buildSnapshot(
 		}
 	}
 
+	switch s.cfg.CloudMigration.ResourceStorageType {
+	case cloudmigration.ResourceStorageTypeDb:
+		if err := s.buildSnapshotWithDBStorage(ctx, snapshotMeta.UID, snapshotWriter, resourcesGroupedByType, maxItemsPerPartition); err != nil {
+			return fmt.Errorf("building snapshot with database storage: %w", err)
+		}
+		s.log.Debug(fmt.Sprintf("buildSnapshot: wrote data partitions with database storage in %d ms", time.Since(start).Milliseconds()))
+
+	case cloudmigration.ResourceStorageTypeFs:
+		if err := s.buildSnapshotWithFSStorage(publicKey[:], metadata, snapshotWriter, resourcesGroupedByType, maxItemsPerPartition); err != nil {
+			return fmt.Errorf("building snapshot with file system storage: %w", err)
+		}
+		s.log.Debug(fmt.Sprintf("buildSnapshot: wrote data partitions with file system storage in %d ms", time.Since(start).Milliseconds()))
+
+	default:
+		return fmt.Errorf("unknown resource storage type, check your configuration and try again: %q", s.cfg.CloudMigration.ResourceStorageType)
+	}
+
+	// update snapshot status to pending upload with retries
+	if err := s.updateSnapshotWithRetries(ctx, cloudmigration.UpdateSnapshotCmd{
+		UID:                    snapshotMeta.UID,
+		SessionID:              snapshotMeta.SessionUID,
+		Status:                 cloudmigration.SnapshotStatusPendingUpload,
+		LocalResourcesToCreate: localSnapshotResource,
+		PublicKey:              publicKey[:],
+	}); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Service) buildSnapshotWithDBStorage(ctx context.Context, snapshotUID string, snapshotWriter *snapshot.SnapshotWriter, resourcesGroupedByType map[cloudmigration.MigrateDataType][]snapshot.MigrateDataRequestItemDTO, maxItemsPerPartition uint32) error {
+	for _, resourceType := range currentMigrationTypes {
+		i := 0
+		for chunk := range slices.Chunk(resourcesGroupedByType[resourceType], int(maxItemsPerPartition)) {
+			encoded, err := snapshotWriter.EncodePartition(chunk)
+			if err != nil {
+				return fmt.Errorf("encoding snapshot partition: %w", err)
+			}
+			if err := s.store.StorePartition(ctx, snapshotUID, string(resourceType), i, encoded); err != nil {
+				return fmt.Errorf("storing partition into database: %w", err)
+			}
+			i += 1
+		}
+	}
+
+	return nil
+}
+
+func (s *Service) buildSnapshotWithFSStorage(publicKey, metadata []byte, snapshotWriter *snapshot.SnapshotWriter, resourcesGroupedByType map[cloudmigration.MigrateDataType][]snapshot.MigrateDataRequestItemDTO, maxItemsPerPartition uint32) error {
 	for _, resourceType := range currentMigrationTypes {
 		for chunk := range slices.Chunk(resourcesGroupedByType[resourceType], int(maxItemsPerPartition)) {
 			if err := snapshotWriter.Write(string(resourceType), chunk); err != nil {
@@ -613,8 +664,6 @@ func (s *Service) buildSnapshot(
 		}
 	}
 
-	s.log.Debug(fmt.Sprintf("buildSnapshot: wrote data files in %d ms", time.Since(start).Milliseconds()))
-
 	// Add the grafana generated public key to the index file so gms can use it to decrypt the snapshot files later.
 	// This works because the snapshot files are being encrypted with
 	// the grafana generated private key + the gms public key.
@@ -625,18 +674,6 @@ func (s *Service) buildSnapshot(
 		return fmt.Errorf("finishing writing snapshot files and generating index file: %w", err)
 	}
 
-	s.log.Debug(fmt.Sprintf("buildSnapshot: finished snapshot in %d ms", time.Since(start).Milliseconds()))
-
-	// update snapshot status to pending upload with retries
-	if err := s.updateSnapshotWithRetries(ctx, cloudmigration.UpdateSnapshotCmd{
-		UID:                    snapshotMeta.UID,
-		SessionID:              snapshotMeta.SessionUID,
-		Status:                 cloudmigration.SnapshotStatusPendingUpload,
-		LocalResourcesToCreate: localSnapshotResource,
-	}); err != nil {
-		return err
-	}
-
 	return nil
 }
 
@@ -654,13 +691,94 @@ func (s *Service) uploadSnapshot(ctx context.Context, session *cloudmigration.Cl
 		s.log.Debug(fmt.Sprintf("uploadSnapshot: method completed in %d ms", time.Since(start).Milliseconds()))
 	}()
 
+	switch s.cfg.CloudMigration.ResourceStorageType {
+	case cloudmigration.ResourceStorageTypeDb:
+		if err := s.uploadSnapshotWithDBStorage(ctx, session, snapshotMeta, uploadUrl); err != nil {
+			return fmt.Errorf("uploading snapshot with database storage: %w", err)
+		}
+
+	case cloudmigration.ResourceStorageTypeFs:
+		if err := s.uploadSnapshotWithFSStorage(ctx, session, snapshotMeta, uploadUrl); err != nil {
+			return fmt.Errorf("uploading snapshot with file system storage: %w", err)
+		}
+
+	default:
+		return fmt.Errorf("unknown resource storage type, check your configuration and try again: %q", s.cfg.CloudMigration.ResourceStorageType)
+	}
+
+	s.log.Info("successfully uploaded snapshot", "snapshotUid", snapshotMeta.UID, "cloud_snapshotUid", snapshotMeta.GMSSnapshotUID)
+
+	// update snapshot status to processing with retries
+	if err := s.updateSnapshotWithRetries(ctx, cloudmigration.UpdateSnapshotCmd{
+		UID:       snapshotMeta.UID,
+		SessionID: snapshotMeta.SessionUID,
+		Status:    cloudmigration.SnapshotStatusProcessing,
+	}); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *Service) uploadSnapshotWithDBStorage(ctx context.Context, session *cloudmigration.CloudMigrationSession, snapshotMeta *cloudmigration.CloudMigrationSnapshot, uploadUrl string) error {
+	index, err := s.store.GetIndex(ctx, session.OrgID, snapshotMeta.SessionUID, snapshotMeta.UID)
+	if err != nil {
+		return fmt.Errorf("fetching index from database: %w", err)
+	}
+
+	snapshotIndex := snapshot.Index{
+		Version:        1,
+		EncryptionAlgo: index.EncryptionAlgo,
+		PublicKey:      index.PublicKey,
+		Metadata:       index.Metadata,
+		Items:          make(map[string][]string),
+	}
+
+	var partitionToFileName = func(resourceType string, partitionNumber int) string {
+		return fmt.Sprintf("%+v_%+v", resourceType, partitionNumber)
+	}
+
+	for resourceType, partitionsNumbers := range index.Items {
+		for _, partitionNumber := range partitionsNumbers {
+			fileName := partitionToFileName(resourceType, partitionNumber)
+			snapshotIndex.Items[resourceType] = append(snapshotIndex.Items[resourceType], fileName)
+
+			key := fmt.Sprintf("%d/snapshots/%s/%+v", session.StackID, snapshotMeta.GMSSnapshotUID, fileName)
+
+			partition, err := s.store.GetPartition(ctx, snapshotMeta.UID, resourceType, partitionNumber)
+			if err != nil {
+				return fmt.Errorf("fetching partition from database: %w", err)
+			}
+			if err = s.objectStorage.PresignedURLUpload(ctx, uploadUrl, key, bytes.NewReader(partition.Data)); err != nil {
+				return fmt.Errorf("uploading file using presigned url: %w", err)
+			}
+		}
+	}
+
+	key := fmt.Sprintf("%d/snapshots/%s/%s", session.StackID, snapshotMeta.GMSSnapshotUID, "index.json")
+
+	buffer, err := snapshot.EncodeIndex(snapshotIndex)
+
+	if err != nil {
+		return fmt.Errorf("encoding snapshot index for upload: %w", err)
+	}
+
+	if err = s.objectStorage.PresignedURLUpload(ctx, uploadUrl, key, bytes.NewReader(buffer)); err != nil {
+		return fmt.Errorf("uploading index file using presigned url: %w", err)
+	}
+
+	return nil
+}
+
+func (s *Service) uploadSnapshotWithFSStorage(ctx context.Context, session *cloudmigration.CloudMigrationSession, snapshotMeta *cloudmigration.CloudMigrationSnapshot, uploadUrl string) error {
 	indexFilePath := filepath.Join(snapshotMeta.LocalDir, "index.json")
+
+	start := time.Now()
 	// LocalDir can be set in the configuration, therefore the file path can be set to any path.
 	// nolint:gosec
 	indexFile, err := os.Open(indexFilePath)
 	if err != nil {
-		// TODO: Clean this notice once we've fixed the HA bug
-		return fmt.Errorf("opening index files: %w. If you are running Grafana in a highly-available setup, try scaling down to one replica to avoid a known bug: https://github.com/grafana/grafana/issues/107264", err)
+		return fmt.Errorf("opening index files: %w. If you are running Grafana in a highly-available setup, try setting cloud_migration.resource_storage_type to 'db' or scaling down to one replica", err)
 	}
 	defer func() {
 		if closeErr := indexFile.Close(); closeErr != nil {
@@ -679,8 +797,6 @@ func (s *Service) uploadSnapshot(ctx context.Context, session *cloudmigration.Cl
 	}
 	readIndexSpan.End()
 
-	s.log.Debug(fmt.Sprintf("uploadSnapshot: read index file in %d ms", time.Since(start).Milliseconds()))
-
 	uploadCtx, uploadSpan := s.tracer.Start(ctx, "CloudMigrationService.uploadSnapshot.uploadDataFiles")
 	// Upload the data files.
 	for _, fileNames := range index.Items {
@@ -724,16 +840,6 @@ func (s *Service) uploadSnapshot(ctx context.Context, session *cloudmigration.Cl
 	uploadSpan.End()
 
 	s.log.Debug(fmt.Sprintf("uploadSnapshot: uploaded index file in %d ms", time.Since(start).Milliseconds()))
-	s.log.Info("successfully uploaded snapshot", "snapshotUid", snapshotMeta.UID, "cloud_snapshotUid", snapshotMeta.GMSSnapshotUID)
-
-	// update snapshot status to processing with retries
-	if err := s.updateSnapshotWithRetries(ctx, cloudmigration.UpdateSnapshotCmd{
-		UID:       snapshotMeta.UID,
-		SessionID: snapshotMeta.SessionUID,
-		Status:    cloudmigration.SnapshotStatusProcessing,
-	}); err != nil {
-		return err
-	}
 
 	return nil
 }

pkg/services/cloudmigration/cloudmigrationimpl/store.go

@@ -12,8 +12,11 @@ type store interface {
 	GetCloudMigrationSessionList(ctx context.Context, orgID int64) ([]*cloudmigration.CloudMigrationSession, error)
 	DeleteMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, []cloudmigration.CloudMigrationSnapshot, error)
 
-	CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) (string, error)
+	CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) error
 	UpdateSnapshot(ctx context.Context, snapshot cloudmigration.UpdateSnapshotCmd) error
+	GetIndex(ctx context.Context, orgID int64, sessionUID string, snapshotUID string) (cloudmigration.CloudMigrationSnapshotIndex, error)
+	GetPartition(ctx context.Context, snapshotUID string, resourceType string, partitionNumber int) (cloudmigration.CloudMigrationSnapshotPartition, error)
+	StorePartition(ctx context.Context, snapshotUID string, resourceType string, partitionNumber int, data []byte) error
 	GetSnapshotByUID(ctx context.Context, orgID int64, sessUid, id string, params cloudmigration.SnapshotResultQueryParams) (*cloudmigration.CloudMigrationSnapshot, error)
 	GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error)
 }

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -121,6 +121,7 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, orgID int64
 		SessionUID: uid,
 		Page:       1,
 		Limit:      GetAllSnapshots,
+		OrgID:      orgID,
 	}
 	snapshots, err := ss.GetSnapshotList(ctx, q)
 	if err != nil {
@@ -129,12 +130,13 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, orgID int64
 
 	err = ss.db.InTransaction(ctx, func(ctx context.Context) error {
 		for _, snapshot := range snapshots {
-			err := ss.deleteSnapshotResources(ctx, snapshot.UID)
-			if err != nil {
+			if err := ss.deleteSnapshotResources(ctx, snapshot.UID); err != nil {
 				return fmt.Errorf("deleting snapshot resource from db: %w", err)
 			}
-			err = ss.deleteSnapshot(ctx, snapshot.UID)
-			if err != nil {
+			if err := ss.deleteSnapshotPartitions(ctx, snapshot.UID); err != nil {
+				return fmt.Errorf("deleting snapshot partitions: %w", err)
+			}
+			if err := ss.deleteSnapshot(ctx, snapshot.UID); err != nil {
 				return fmt.Errorf("deleting snapshot from db: %w", err)
 			}
 		}
@@ -166,33 +168,33 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, orgID int64
 	return &c, snapshots, nil
 }
 
-func (ss *sqlStore) CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) (string, error) {
+func (ss *sqlStore) CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) error {
 	if snapshot.SessionUID == "" {
-		return "", fmt.Errorf("sessionUID is required")
+		return fmt.Errorf("sessionUID is required")
 	}
-
 	if snapshot.UID == "" {
-		snapshot.UID = util.GenerateShortUID()
+		return fmt.Errorf("snapshot uid is required")
 	}
 
-	if err := ss.secretsStore.Set(ctx, secretskv.AllOrganizations, snapshot.UID, secretType, string(snapshot.EncryptionKey)); err != nil {
-		return "", err
+	if err := ss.secretsStore.Set(ctx, secretskv.AllOrganizations, snapshot.UID, secretType, string(snapshot.GMSPublicKey)); err != nil {
+		return err
 	}
 
 	err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		snapshot.Created = time.Now()
 		snapshot.Updated = time.Now()
 
-		_, err := sess.Insert(&snapshot)
+		_, err := sess.InsertOne(&snapshot)
 		if err != nil {
 			return err
 		}
+
 		return nil
 	})
 	if err != nil {
-		return "", err
+		return err
 	}
-	return snapshot.UID, nil
+	return nil
 }
 
 // UpdateSnapshot takes a command containing a snapshot uid and any updates to apply to the snapshot.
@@ -232,19 +234,133 @@ func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.Up
 			return err
 		}
 	}
+	if update.PublicKey != nil {
+		if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			rawSQL := "UPDATE cloud_migration_snapshot SET public_key=? WHERE session_uid=? AND uid=?"
+			if _, err := sess.Exec(rawSQL, update.PublicKey, update.SessionID, update.UID); err != nil {
+				return fmt.Errorf("updating snapshot public key for uid %s: %w", update.UID, err)
+			}
+			return nil
+		}); err != nil {
+			return err
+		}
+	}
 
 	return nil
 }
 
-func (ss *sqlStore) deleteSnapshot(ctx context.Context, snapshotUid string) error {
-	return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
-		_, err := sess.Delete(cloudmigration.CloudMigrationSnapshot{
-			UID: snapshotUid,
+func (ss *sqlStore) StorePartition(ctx context.Context, snapshotUID string, resourceType string, partitionNumber int, data []byte) error {
+	return ss.db.InTransaction(ctx, func(ctx context.Context) error {
+		return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			_, err := sess.Insert(cloudmigration.CloudMigrationSnapshotPartition{
+				SnapshotUID:     snapshotUID,
+				ResourceType:    resourceType,
+				PartitionNumber: partitionNumber,
+				Data:            data,
+			})
+			if err != nil {
+				return fmt.Errorf("inserting snapshot partition into database: %w", err)
+			}
+			return nil
 		})
-		return err
 	})
 }
 
+func (ss *sqlStore) GetIndex(ctx context.Context, orgID int64, sessionUID string, snapshotUID string) (cloudmigration.CloudMigrationSnapshotIndex, error) {
+	var snap *cloudmigration.CloudMigrationSnapshot
+	partitions := make([]cloudmigration.CloudMigrationSnapshotPartition, 0)
+
+	if err := ss.db.InTransaction(ctx, func(ctx context.Context) error {
+		return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			s, err := ss.getSnapshotByUID(ctx, orgID, sessionUID, snapshotUID)
+			if err != nil {
+				return fmt.Errorf("fetching snapshot from database: %w", err)
+			}
+
+			snap = s
+			if err := sess.OrderBy("cloud_migration_snapshot_partition.resource_type,cloud_migration_snapshot_partition.partition_number ASC").Find(&partitions, &cloudmigration.CloudMigrationSnapshotPartition{SnapshotUID: snapshotUID}); err != nil {
+				return fmt.Errorf("fetching partition from database: %w", err)
+			}
+			if secret, found, err := ss.secretsStore.Get(ctx, secretskv.AllOrganizations, snap.UID, secretType); err != nil {
+				return err
+			} else if !found {
+				return fmt.Errorf("encryption key not found for snapshot with UID %s", snap.UID)
+			} else {
+				snap.GMSPublicKey = []byte(secret)
+			}
+
+			return nil
+		})
+	}); err != nil {
+		return cloudmigration.CloudMigrationSnapshotIndex{}, err
+	}
+
+	partitionsByResourceType := make(map[string][]int)
+	for _, partition := range partitions {
+		partitionsByResourceType[partition.ResourceType] = append(partitionsByResourceType[partition.ResourceType], partition.PartitionNumber)
+	}
+
+	return cloudmigration.CloudMigrationSnapshotIndex{
+		EncryptionAlgo: snap.EncryptionAlgo,
+		PublicKey:      snap.PublicKey,
+		Metadata:       snap.Metadata,
+		Items:          partitionsByResourceType,
+	}, nil
+}
+
+func (ss *sqlStore) GetPartition(ctx context.Context, snapshotUID string, resourceType string, partitionNumber int) (cloudmigration.CloudMigrationSnapshotPartition, error) {
+	var partition cloudmigration.CloudMigrationSnapshotPartition
+
+	err := ss.db.InTransaction(ctx, func(ctx context.Context) error {
+		return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			if _, err := sess.Where("snapshot_uid = ? AND resource_type = ?  AND partition_number = ?", snapshotUID, resourceType, partitionNumber).Get(&partition); err != nil {
+				return fmt.Errorf("fetching partition from database: %w", err)
+			}
+			return nil
+		})
+	})
+
+	return partition, err
+}
+
+func (ss *sqlStore) deleteSnapshot(ctx context.Context, snapshotUid string) error {
+	return ss.db.InTransaction(ctx, func(ctx context.Context) error {
+		return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			if _, err := sess.Delete(cloudmigration.CloudMigrationSnapshot{
+				UID: snapshotUid,
+			}); err != nil {
+				return fmt.Errorf("deleting snapshot: %w", err)
+			}
+			return nil
+		})
+	})
+}
+
+func (ss *sqlStore) getSnapshotByUID(ctx context.Context, orgID int64, sessionUID string, snapshotUID string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	session, err := ss.GetMigrationSessionByUID(ctx, orgID, sessionUID)
+	if err != nil || session == nil {
+		return nil, err
+	}
+
+	// now we get the snapshot
+	var snapshot cloudmigration.CloudMigrationSnapshot
+	err = ss.db.WithDbSession(ctx, func(sess *db.Session) error {
+		exist, err := sess.Where("session_uid=? AND uid=?", sessionUID, snapshotUID).Get(&snapshot)
+		if err != nil {
+			return err
+		}
+		if !exist {
+			return cloudmigration.ErrSnapshotNotFound
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &snapshot, nil
+}
+
 func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, orgID int64, sessionUid, uid string, params cloudmigration.SnapshotResultQueryParams) (*cloudmigration.CloudMigrationSnapshot, error) {
 	// first we check if the session exists, using orgId and sessionUid
 	session, err := ss.GetMigrationSessionByUID(ctx, orgID, sessionUid)
@@ -273,7 +389,7 @@ func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, orgID int64, sessionUi
 	} else if !found {
 		return &snapshot, fmt.Errorf("encryption key not found for snapshot with UID %s", snapshot.UID)
 	} else {
-		snapshot.EncryptionKey = []byte(secret)
+		snapshot.GMSPublicKey = []byte(secret)
 	}
 
 	resources, err := ss.getSnapshotResources(ctx, uid, params)
@@ -291,6 +407,12 @@ func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, orgID int64, sessionUi
 // GetSnapshotList returns snapshots without resources included. Use GetSnapshotByUID to get individual snapshot results.
 // passing GetAllSnapshots will return all the elements regardless of the page
 func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
+	if query.OrgID == 0 {
+		return nil, fmt.Errorf("org id is required")
+	}
+	if query.SessionUID == "" {
+		return nil, fmt.Errorf("session uid is required")
+	}
 	var snapshots = make([]cloudmigration.CloudMigrationSnapshot, 0)
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
 		sess.Join("INNER", "cloud_migration_session",
@@ -310,13 +432,14 @@ func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.Li
 	if err != nil {
 		return nil, err
 	}
+
 	for i, snapshot := range snapshots {
 		if secret, found, err := ss.secretsStore.Get(ctx, secretskv.AllOrganizations, snapshot.UID, secretType); err != nil {
 			return nil, err
 		} else if !found {
 			return nil, fmt.Errorf("encryption key not found for snapshot with UID %s", snapshot.UID)
 		} else {
-			snapshot.EncryptionKey = []byte(secret)
+			snapshot.GMSPublicKey = []byte(secret)
 		}
 
 		if stats, err := ss.getSnapshotResourceStats(ctx, snapshot.UID); err != nil {
@@ -531,6 +654,17 @@ func (ss *sqlStore) deleteSnapshotResources(ctx context.Context, snapshotUid str
 	})
 }
 
+func (ss *sqlStore) deleteSnapshotPartitions(ctx context.Context, snapshotUid string) error {
+	return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		if _, err := sess.Delete(cloudmigration.CloudMigrationSnapshotPartition{
+			SnapshotUID: snapshotUid,
+		}); err != nil {
+			return fmt.Errorf("deleting snapshot partitions: %w", err)
+		}
+		return nil
+	})
+}
+
 func (ss *sqlStore) encryptToken(ctx context.Context, cm *cloudmigration.CloudMigrationSession) error {
 	s, err := ss.secretsService.Encrypt(ctx, []byte(cm.AuthToken), secrets.WithoutScope())
 	if err != nil {

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store_test.go

@@ -1,13 +1,18 @@
 package cloudmigrationimpl
 
 import (
+	"bytes"
 	"context"
+	cryptoRand "crypto/rand"
 	"encoding/base64"
 	"fmt"
 	"strconv"
 	"testing"
 
 	"github.com/google/uuid"
+	snapshot "github.com/grafana/grafana-cloud-migration-snapshot/src"
+	"github.com/grafana/grafana-cloud-migration-snapshot/src/contracts"
+	"github.com/grafana/grafana-cloud-migration-snapshot/src/infra/crypto"
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
 	fakeSecrets "github.com/grafana/grafana/pkg/services/secrets/fakes"
@@ -15,6 +20,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/nacl/box"
 )
 
 func Test_GetAllCloudMigrationSessions(t *testing.T) {
@@ -112,17 +118,18 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NoError(t, err)
 
 		// create a snapshot
+		uid := uuid.NewString()
 		cmr := cloudmigration.CloudMigrationSnapshot{
+			UID:        uid,
 			SessionUID: session.UID,
 			Status:     cloudmigration.SnapshotStatusCreating,
 		}
 
-		snapshotUid, err := s.CreateSnapshot(ctx, cmr)
+		err = s.CreateSnapshot(ctx, cmr)
 		require.NoError(t, err)
-		require.NotEmpty(t, snapshotUid)
 
 		//retrieve it from the db
-		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, cloudmigration.SnapshotResultQueryParams{
+		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, uid, cloudmigration.SnapshotResultQueryParams{
 			ResultPage:  1,
 			ResultLimit: 100,
 			SortColumn:  cloudmigration.SortColumnID,
@@ -132,11 +139,11 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.Equal(t, cloudmigration.SnapshotStatusCreating, snapshot.Status)
 
 		// update its status
-		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{UID: snapshotUid, Status: cloudmigration.SnapshotStatusCreating, SessionID: session.UID})
+		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{UID: uid, Status: cloudmigration.SnapshotStatusCreating, SessionID: session.UID})
 		require.NoError(t, err)
 
 		//retrieve it again
-		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, cloudmigration.SnapshotResultQueryParams{
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, uid, cloudmigration.SnapshotResultQueryParams{
 			ResultPage:  1,
 			ResultLimit: 100,
 			SortColumn:  cloudmigration.SortColumnID,
@@ -152,11 +159,11 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.Equal(t, *snapshot, snapshots[0])
 
 		// delete snapshot
-		err = s.deleteSnapshot(ctx, snapshotUid)
+		err = s.deleteSnapshot(ctx, uid)
 		require.NoError(t, err)
 
 		// now we expect not to find the snapshot
-		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, cloudmigration.SnapshotResultQueryParams{
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, uid, cloudmigration.SnapshotResultQueryParams{
 			ResultPage:  1,
 			ResultLimit: 100,
 			SortColumn:  cloudmigration.SortColumnID,
@@ -174,12 +181,13 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NoError(t, err)
 
 		// create a snapshot
-		snapshotUid, err := s.CreateSnapshot(ctx, cloudmigration.CloudMigrationSnapshot{
+		uid := uuid.NewString()
+		err = s.CreateSnapshot(ctx, cloudmigration.CloudMigrationSnapshot{
+			UID:        uid,
 			SessionUID: session.UID,
 			Status:     cloudmigration.SnapshotStatusCreating,
 		})
 		require.NoError(t, err)
-		require.NotEmpty(t, snapshotUid)
 
 		// Generate 50,001 test resources in order to test both update conditions (reached the batch limit or reached the end)
 		const numResources = 50001
@@ -196,7 +204,7 @@ func Test_SnapshotManagement(t *testing.T) {
 
 		// Update the snapshot with the resources to create
 		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
-			UID:                    snapshotUid,
+			UID:                    uid,
 			Status:                 cloudmigration.SnapshotStatusPendingUpload,
 			SessionID:              session.UID,
 			LocalResourcesToCreate: resources,
@@ -204,7 +212,7 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NoError(t, err)
 
 		// Get the Snapshot and ensure it's in the right state
-		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, cloudmigration.SnapshotResultQueryParams{
+		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, uid, cloudmigration.SnapshotResultQueryParams{
 			ResultPage:  1,
 			ResultLimit: numResources,
 			SortColumn:  cloudmigration.SortColumnID,
@@ -226,7 +234,7 @@ func Test_SnapshotManagement(t *testing.T) {
 
 		// Update the snapshot with the resources to update
 		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
-			UID:                    snapshotUid,
+			UID:                    uid,
 			Status:                 cloudmigration.SnapshotStatusFinished,
 			SessionID:              session.UID,
 			CloudResourcesToUpdate: snapshot.Resources,
@@ -234,7 +242,7 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NoError(t, err)
 
 		// Get the Snapshot and ensure it's in the right state
-		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, cloudmigration.SnapshotResultQueryParams{
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, uid, cloudmigration.SnapshotResultQueryParams{
 			ResultPage:  1,
 			ResultLimit: numResources,
 			SortColumn:  cloudmigration.SortColumnID,
@@ -637,7 +645,7 @@ func TestGetSnapshotList(t *testing.T) {
 	})
 
 	t.Run("return no snapshots if limit is set to 0", func(t *testing.T) {
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, Page: 1, Limit: 0})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, OrgID: 1, Page: 1, Limit: 0})
 		require.NoError(t, err)
 		assert.Empty(t, snapshots)
 	})
@@ -669,7 +677,7 @@ func TestGetSnapshotList(t *testing.T) {
 	})
 
 	t.Run("only the snapshots that belong to a specific session are returned", func(t *testing.T) {
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: "session-uid-that-doesnt-exist", Page: 1, Limit: 100})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: "session-uid-that-doesnt-exist", OrgID: 1, Page: 1, Limit: 100})
 		require.NoError(t, err)
 		assert.Empty(t, snapshots)
 	})
@@ -680,7 +688,7 @@ func TestGetSnapshotList(t *testing.T) {
 		require.NoError(t, err)
 
 		// Fetch the snapshots that belong to the deleted session.
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, Page: 1, Limit: 100})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, OrgID: 1, Page: 1, Limit: 100})
 		require.NoError(t, err)
 
 		// No snapshots should be returned because the session that
@@ -801,3 +809,42 @@ func setUpTest(t *testing.T) (*sqlstore.SQLStore, *sqlStore) {
 func encodeToken(t string) string {
 	return base64.StdEncoding.EncodeToString([]byte(t))
 }
+
+func TestEncodeDecode(t *testing.T) {
+	gmsPublicKey, gmsPrivateKey, err := box.GenerateKey(cryptoRand.Reader)
+	require.NoError(t, err)
+
+	grafanaPublicKey, grafanaPrivateKey, err := box.GenerateKey(cryptoRand.Reader)
+	require.NoError(t, err)
+
+	snapshotWriter, err := snapshot.NewSnapshotWriter(contracts.AssymetricKeys{
+		Public:  gmsPublicKey[:],
+		Private: grafanaPrivateKey[:],
+	},
+		crypto.NewNacl(),
+		"",
+	)
+	require.NoError(t, err)
+
+	chunk := []snapshot.MigrateDataRequestItemDTO{{
+		Type:  snapshot.AlertRuleGroupType,
+		RefID: "foo",
+		Name:  "name",
+		Data:  map[string]any{"a": "b"},
+	}}
+	encoded, err := snapshotWriter.EncodePartition(chunk)
+	require.NoError(t, err)
+
+	require.NoError(t, snapshotWriter.Write("RESOURCE_TYPE", chunk))
+
+	reader := snapshot.NewSnapshotReader(contracts.AssymetricKeys{
+		Public:  grafanaPublicKey[:],
+		Private: gmsPrivateKey[:],
+	},
+		crypto.NewNacl())
+
+	partition, err := reader.ReadFile(bytes.NewReader(encoded))
+	require.NoError(t, err)
+
+	require.Equal(t, chunk, partition.Items)
+}