CloudMigrations: Implement snapshot management apis (#89296)

* add new apis

* add payloads

* create snapshot status type

* add some impl

* finish implementing update

* start implementing build snapshot func

* add more fake build logic

* add cancel endpoint. do some cleanup

* implement GetSnapshot

* implement upload snapshot

* merge onprem status with gms result

* get it working

* update comment

* rename list endpoint

* add query limit and offset

* add helper method to snapshot

* little bit of cleanup

* work on swagger annotations

* manual merge

* generate swagger specs

* clean up curl commands

* fix bugs found during final testing

* fix linter issue

* fix unit test

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go

@@ -7,6 +7,9 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"os"
+	"path/filepath"
+	"sync"
 	"time"
 
 	"github.com/grafana/grafana/pkg/api/response"
@@ -17,7 +20,6 @@ import (
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
 	"github.com/grafana/grafana/pkg/services/cloudmigration/api"
 	"github.com/grafana/grafana/pkg/services/cloudmigration/gmsclient"
-	"github.com/grafana/grafana/pkg/services/contexthandler"
 	"github.com/grafana/grafana/pkg/services/dashboards"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
@@ -25,6 +27,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/gcom"
 	"github.com/grafana/grafana/pkg/services/secrets"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/util"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
@@ -37,6 +40,9 @@ type Service struct {
 	log *log.ConcreteLogger
 	cfg *setting.Cfg
 
+	buildSnapshotMutex sync.Mutex
+	buildSnapshotError bool
+
 	features  featuremgmt.FeatureToggles
 	gmsClient gmsclient.Client
 
@@ -398,7 +404,6 @@ func (s *Service) RunMigration(ctx context.Context, uid string) (*cloudmigration
 		return nil, fmt.Errorf("migrate data error: %w", err)
 	}
 
-	// TODO update cloud migration run schema to treat the result as a first-class citizen
 	respData, err := json.Marshal(resp)
 	if err != nil {
 		s.log.Error("error marshalling migration response data: %w", err)
@@ -419,135 +424,6 @@ func (s *Service) RunMigration(ctx context.Context, uid string) (*cloudmigration
 	return resp, nil
 }
 
-func (s *Service) getMigrationDataJSON(ctx context.Context) (*cloudmigration.MigrateDataRequest, error) {
-	// Data sources
-	dataSources, err := s.getDataSources(ctx)
-	if err != nil {
-		s.log.Error("Failed to get datasources", "err", err)
-		return nil, err
-	}
-
-	// Dashboards
-	dashboards, err := s.getDashboards(ctx)
-	if err != nil {
-		s.log.Error("Failed to get dashboards", "err", err)
-		return nil, err
-	}
-
-	// Folders
-	folders, err := s.getFolders(ctx)
-	if err != nil {
-		s.log.Error("Failed to get folders", "err", err)
-		return nil, err
-	}
-
-	migrationDataSlice := make(
-		[]cloudmigration.MigrateDataRequestItem, 0,
-		len(dataSources)+len(dashboards)+len(folders),
-	)
-	for _, ds := range dataSources {
-		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
-			Type:  cloudmigration.DatasourceDataType,
-			RefID: ds.UID,
-			Name:  ds.Name,
-			Data:  ds,
-		})
-	}
-
-	for _, dashboard := range dashboards {
-		dashboard.Data.Del("id")
-		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
-			Type:  cloudmigration.DashboardDataType,
-			RefID: dashboard.UID,
-			Name:  dashboard.Title,
-			Data:  map[string]any{"dashboard": dashboard.Data},
-		})
-	}
-
-	for _, f := range folders {
-		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
-			Type:  cloudmigration.FolderDataType,
-			RefID: f.UID,
-			Name:  f.Title,
-			Data:  f,
-		})
-	}
-
-	migrationData := &cloudmigration.MigrateDataRequest{
-		Items: migrationDataSlice,
-	}
-
-	return migrationData, nil
-}
-
-func (s *Service) getDataSources(ctx context.Context) ([]datasources.AddDataSourceCommand, error) {
-	dataSources, err := s.dsService.GetAllDataSources(ctx, &datasources.GetAllDataSourcesQuery{})
-	if err != nil {
-		s.log.Error("Failed to get all datasources", "err", err)
-		return nil, err
-	}
-
-	result := []datasources.AddDataSourceCommand{}
-	for _, dataSource := range dataSources {
-		// Decrypt secure json to send raw credentials
-		decryptedData, err := s.secretsService.DecryptJsonData(ctx, dataSource.SecureJsonData)
-		if err != nil {
-			s.log.Error("Failed to decrypt secure json data", "err", err)
-			return nil, err
-		}
-		dataSourceCmd := datasources.AddDataSourceCommand{
-			OrgID:           dataSource.OrgID,
-			Name:            dataSource.Name,
-			Type:            dataSource.Type,
-			Access:          dataSource.Access,
-			URL:             dataSource.URL,
-			User:            dataSource.User,
-			Database:        dataSource.Database,
-			BasicAuth:       dataSource.BasicAuth,
-			BasicAuthUser:   dataSource.BasicAuthUser,
-			WithCredentials: dataSource.WithCredentials,
-			IsDefault:       dataSource.IsDefault,
-			JsonData:        dataSource.JsonData,
-			SecureJsonData:  decryptedData,
-			ReadOnly:        dataSource.ReadOnly,
-			UID:             dataSource.UID,
-		}
-		result = append(result, dataSourceCmd)
-	}
-	return result, err
-}
-
-func (s *Service) getFolders(ctx context.Context) ([]folder.Folder, error) {
-	reqCtx := contexthandler.FromContext(ctx)
-	folders, err := s.folderService.GetFolders(ctx, folder.GetFoldersQuery{
-		SignedInUser: reqCtx.SignedInUser,
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	result := make([]folder.Folder, len(folders))
-	for i, folder := range folders {
-		result[i] = *folder
-	}
-
-	return result, nil
-}
-
-func (s *Service) getDashboards(ctx context.Context) ([]dashboards.Dashboard, error) {
-	dashs, err := s.dashboardService.GetAllDashboards(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	result := make([]dashboards.Dashboard, len(dashs))
-	for i, dashboard := range dashs {
-		result[i] = *dashboard
-	}
-
-	return result, nil
-}
-
 func (s *Service) createMigrationRun(ctx context.Context, cmr cloudmigration.CloudMigrationSnapshot) (string, error) {
 	uid, err := s.store.CreateMigrationRun(ctx, cmr)
 	if err != nil {
@@ -565,13 +441,13 @@ func (s *Service) GetMigrationStatus(ctx context.Context, runUID string) (*cloud
 	return cmr, nil
 }
 
-func (s *Service) GetMigrationRunList(ctx context.Context, migUID string) (*cloudmigration.SnapshotList, error) {
+func (s *Service) GetMigrationRunList(ctx context.Context, migUID string) (*cloudmigration.CloudMigrationRunList, error) {
 	runs, err := s.store.GetMigrationStatusList(ctx, migUID)
 	if err != nil {
 		return nil, fmt.Errorf("retrieving migration statuses from db: %w", err)
 	}
 
-	runList := &cloudmigration.SnapshotList{Runs: []cloudmigration.MigrateDataResponseList{}}
+	runList := &cloudmigration.CloudMigrationRunList{Runs: []cloudmigration.MigrateDataResponseList{}}
 	for _, s := range runs {
 		runList.Runs = append(runList.Runs, cloudmigration.MigrateDataResponseList{
 			RunUID: s.UID,
@@ -589,6 +465,123 @@ func (s *Service) DeleteSession(ctx context.Context, uid string) (*cloudmigratio
 	return c, nil
 }
 
+func (s *Service) CreateSnapshot(ctx context.Context, sessionUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.CreateSnapshot")
+	defer span.End()
+
+	// fetch session for the gms auth token
+	session, err := s.store.GetMigrationSessionByUID(ctx, sessionUid)
+	if err != nil {
+		return nil, fmt.Errorf("fetching migration session for uid %s: %w", sessionUid, err)
+	}
+
+	// query gms to establish new snapshot
+	initResp, err := s.gmsClient.InitializeSnapshot(ctx, *session)
+	if err != nil {
+		return nil, fmt.Errorf("initializing snapshot with GMS for session %s: %w", sessionUid, err)
+	}
+
+	// create new directory for snapshot writing
+	snapshotUid := util.GenerateShortUID()
+	dir := filepath.Join("cloudmigration.snapshots", fmt.Sprintf("snapshot-%s-%s", snapshotUid, initResp.GMSSnapshotUID))
+	err = os.MkdirAll(dir, 0750)
+	if err != nil {
+		return nil, fmt.Errorf("creating snapshot directory: %w", err)
+	}
+
+	// save snapshot to the db
+	snapshot := cloudmigration.CloudMigrationSnapshot{
+		UID:            snapshotUid,
+		SessionUID:     sessionUid,
+		Status:         cloudmigration.SnapshotStatusInitializing,
+		EncryptionKey:  initResp.EncryptionKey,
+		UploadURL:      initResp.UploadURL,
+		GMSSnapshotUID: initResp.GMSSnapshotUID,
+		LocalDir:       dir,
+	}
+
+	uid, err := s.store.CreateSnapshot(ctx, snapshot)
+	if err != nil {
+		return nil, fmt.Errorf("saving snapshot: %w", err)
+	}
+	snapshot.UID = uid
+
+	// start building the snapshot asynchronously while we return a success response to the client
+	go s.buildSnapshot(context.Background(), snapshot)
+
+	return &snapshot, nil
+}
+
+// GetSnapshot returns the on-prem version of a snapshot, supplemented with processing status from GMS
+func (s *Service) GetSnapshot(ctx context.Context, sessionUid string, snapshotUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetSnapshot")
+	defer span.End()
+
+	snapshot, err := s.store.GetSnapshotByUID(ctx, snapshotUid)
+	if err != nil {
+		return nil, fmt.Errorf("fetching snapshot for uid %s: %w", snapshotUid, err)
+	}
+
+	session, err := s.store.GetMigrationSessionByUID(ctx, sessionUid)
+	if err != nil {
+		return nil, fmt.Errorf("fetching session for uid %s: %w", sessionUid, err)
+	}
+
+	if snapshot.ShouldQueryGMS() {
+		// ask GMS for status if it's in the cloud
+		snapshotMeta, err := s.gmsClient.GetSnapshotStatus(ctx, *session, *snapshot)
+		if err != nil {
+			return nil, fmt.Errorf("error fetching snapshot status from GMS: sessionUid: %s, snapshotUid: %s", sessionUid, snapshotUid)
+		}
+
+		// grab any result available
+		// TODO: figure out a more intelligent way to do this, will depend on GMS apis
+		snapshot.Result = snapshotMeta.Result
+
+		if snapshotMeta.Status == cloudmigration.SnapshotStatusFinished {
+			// we need to update the snapshot in our db before reporting anything finished to the client
+			if err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+				UID:    snapshot.UID,
+				Status: cloudmigration.SnapshotStatusFinished,
+				Result: snapshot.Result,
+			}); err != nil {
+				return nil, fmt.Errorf("error updating snapshot status: %w", err)
+			}
+		}
+	}
+
+	return snapshot, nil
+}
+
+func (s *Service) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
+	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetSnapshotList")
+	defer span.End()
+
+	snapshotList, err := s.store.GetSnapshotList(ctx, query)
+	if err != nil {
+		return nil, fmt.Errorf("fetching snapshots for session uid %s: %w", query.SessionUID, err)
+	}
+	return snapshotList, nil
+}
+
+func (s *Service) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.UploadSnapshot")
+	defer span.End()
+
+	snapshot, err := s.GetSnapshot(ctx, sessionUid, snapshotUid)
+	if err != nil {
+		return fmt.Errorf("fetching snapshot with uid %s: %w", snapshotUid, err)
+	}
+
+	s.log.Info("Uploading snapshot with GMS ID %s in local directory %s to url %s", snapshot.GMSSnapshotUID, snapshot.LocalDir, snapshot.UploadURL)
+	s.log.Debug("UploadSnapshot not yet implemented, faking it")
+
+	// start uploading the snapshot asynchronously while we return a success response to the client
+	go s.uploadSnapshot(context.Background(), *snapshot)
+
+	return nil
+}
+
 func (s *Service) parseCloudMigrationConfig() (string, error) {
 	if s.cfg == nil {
 		return "", fmt.Errorf("cfg cannot be nil")

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration_noop.go

@@ -44,7 +44,7 @@ func (s *NoopServiceImpl) GetMigrationStatus(ctx context.Context, runUID string)
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
-func (s *NoopServiceImpl) GetMigrationRunList(ctx context.Context, uid string) (*cloudmigration.SnapshotList, error) {
+func (s *NoopServiceImpl) GetMigrationRunList(ctx context.Context, uid string) (*cloudmigration.CloudMigrationRunList, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
@@ -59,3 +59,19 @@ func (s *NoopServiceImpl) CreateMigrationRun(context.Context, cloudmigration.Clo
 func (s *NoopServiceImpl) RunMigration(context.Context, string) (*cloudmigration.MigrateDataResponse, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
+
+func (s *NoopServiceImpl) CreateSnapshot(ctx context.Context, sessionUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	return nil, cloudmigration.ErrFeatureDisabledError
+}
+
+func (s *NoopServiceImpl) GetSnapshot(ctx context.Context, sessionUid string, snapshotUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	return nil, cloudmigration.ErrFeatureDisabledError
+}
+
+func (s *NoopServiceImpl) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
+	return nil, cloudmigration.ErrFeatureDisabledError
+}
+
+func (s *NoopServiceImpl) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+	return cloudmigration.ErrFeatureDisabledError
+}

pkg/services/cloudmigration/cloudmigrationimpl/fake/cloudmigration_fake.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
 	"github.com/grafana/grafana/pkg/services/gcom"
+	"github.com/grafana/grafana/pkg/util"
 )
 
 var fixedDate = time.Date(2024, 6, 5, 17, 30, 40, 0, time.UTC)
@@ -122,14 +123,61 @@ func (m FakeServiceImpl) GetMigrationStatus(_ context.Context, _ string) (*cloud
 	}, nil
 }
 
-func (m FakeServiceImpl) GetMigrationRunList(_ context.Context, _ string) (*cloudmigration.SnapshotList, error) {
+func (m FakeServiceImpl) GetMigrationRunList(_ context.Context, _ string) (*cloudmigration.CloudMigrationRunList, error) {
 	if m.ReturnError {
 		return nil, fmt.Errorf("mock error")
 	}
-	return &cloudmigration.SnapshotList{
+	return &cloudmigration.CloudMigrationRunList{
 		Runs: []cloudmigration.MigrateDataResponseList{
 			{RunUID: "fake_run_uid_1"},
 			{RunUID: "fake_run_uid_2"},
 		},
 	}, nil
 }
+
+func (m FakeServiceImpl) CreateSnapshot(ctx context.Context, sessionUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	if m.ReturnError {
+		return nil, fmt.Errorf("mock error")
+	}
+	return &cloudmigration.CloudMigrationSnapshot{
+		UID:        util.GenerateShortUID(),
+		SessionUID: sessionUid,
+		Status:     cloudmigration.SnapshotStatusUnknown,
+	}, nil
+}
+
+func (m FakeServiceImpl) GetSnapshot(ctx context.Context, sessionUid string, snapshotUid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	if m.ReturnError {
+		return nil, fmt.Errorf("mock error")
+	}
+	return &cloudmigration.CloudMigrationSnapshot{
+		UID:        util.GenerateShortUID(),
+		SessionUID: sessionUid,
+		Status:     cloudmigration.SnapshotStatusUnknown,
+	}, nil
+}
+
+func (m FakeServiceImpl) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
+	if m.ReturnError {
+		return nil, fmt.Errorf("mock error")
+	}
+	return []cloudmigration.CloudMigrationSnapshot{
+		{
+			UID:        util.GenerateShortUID(),
+			SessionUID: query.SessionUID,
+			Status:     cloudmigration.SnapshotStatusUnknown,
+		},
+		{
+			UID:        util.GenerateShortUID(),
+			SessionUID: query.SessionUID,
+			Status:     cloudmigration.SnapshotStatusUnknown,
+		},
+	}, nil
+}
+
+func (m FakeServiceImpl) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+	if m.ReturnError {
+		return fmt.Errorf("mock error")
+	}
+	return nil
+}

pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

@@ -0,0 +1,234 @@
+package cloudmigrationimpl
+
+import (
+	"context"
+	"time"
+
+	"github.com/grafana/grafana/pkg/services/cloudmigration"
+	"github.com/grafana/grafana/pkg/services/contexthandler"
+	"github.com/grafana/grafana/pkg/services/dashboards"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/folder"
+	"github.com/grafana/grafana/pkg/util/retryer"
+)
+
+func (s *Service) getMigrationDataJSON(ctx context.Context) (*cloudmigration.MigrateDataRequest, error) {
+	// Data sources
+	dataSources, err := s.getDataSources(ctx)
+	if err != nil {
+		s.log.Error("Failed to get datasources", "err", err)
+		return nil, err
+	}
+
+	// Dashboards
+	dashboards, err := s.getDashboards(ctx)
+	if err != nil {
+		s.log.Error("Failed to get dashboards", "err", err)
+		return nil, err
+	}
+
+	// Folders
+	folders, err := s.getFolders(ctx)
+	if err != nil {
+		s.log.Error("Failed to get folders", "err", err)
+		return nil, err
+	}
+
+	migrationDataSlice := make(
+		[]cloudmigration.MigrateDataRequestItem, 0,
+		len(dataSources)+len(dashboards)+len(folders),
+	)
+
+	for _, ds := range dataSources {
+		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
+			Type:  cloudmigration.DatasourceDataType,
+			RefID: ds.UID,
+			Name:  ds.Name,
+			Data:  ds,
+		})
+	}
+
+	for _, dashboard := range dashboards {
+		dashboard.Data.Del("id")
+		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
+			Type:  cloudmigration.DashboardDataType,
+			RefID: dashboard.UID,
+			Name:  dashboard.Title,
+			Data:  map[string]any{"dashboard": dashboard.Data},
+		})
+	}
+
+	for _, f := range folders {
+		migrationDataSlice = append(migrationDataSlice, cloudmigration.MigrateDataRequestItem{
+			Type:  cloudmigration.FolderDataType,
+			RefID: f.UID,
+			Name:  f.Title,
+			Data:  f,
+		})
+	}
+
+	migrationData := &cloudmigration.MigrateDataRequest{
+		Items: migrationDataSlice,
+	}
+
+	return migrationData, nil
+}
+
+func (s *Service) getDataSources(ctx context.Context) ([]datasources.AddDataSourceCommand, error) {
+	dataSources, err := s.dsService.GetAllDataSources(ctx, &datasources.GetAllDataSourcesQuery{})
+	if err != nil {
+		s.log.Error("Failed to get all datasources", "err", err)
+		return nil, err
+	}
+
+	result := []datasources.AddDataSourceCommand{}
+	for _, dataSource := range dataSources {
+		// Decrypt secure json to send raw credentials
+		decryptedData, err := s.secretsService.DecryptJsonData(ctx, dataSource.SecureJsonData)
+		if err != nil {
+			s.log.Error("Failed to decrypt secure json data", "err", err)
+			return nil, err
+		}
+		dataSourceCmd := datasources.AddDataSourceCommand{
+			OrgID:           dataSource.OrgID,
+			Name:            dataSource.Name,
+			Type:            dataSource.Type,
+			Access:          dataSource.Access,
+			URL:             dataSource.URL,
+			User:            dataSource.User,
+			Database:        dataSource.Database,
+			BasicAuth:       dataSource.BasicAuth,
+			BasicAuthUser:   dataSource.BasicAuthUser,
+			WithCredentials: dataSource.WithCredentials,
+			IsDefault:       dataSource.IsDefault,
+			JsonData:        dataSource.JsonData,
+			SecureJsonData:  decryptedData,
+			ReadOnly:        dataSource.ReadOnly,
+			UID:             dataSource.UID,
+		}
+		result = append(result, dataSourceCmd)
+	}
+	return result, err
+}
+
+func (s *Service) getFolders(ctx context.Context) ([]folder.Folder, error) {
+	reqCtx := contexthandler.FromContext(ctx)
+	folders, err := s.folderService.GetFolders(ctx, folder.GetFoldersQuery{
+		SignedInUser: reqCtx.SignedInUser,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	result := make([]folder.Folder, len(folders))
+	for i, folder := range folders {
+		result[i] = *folder
+	}
+
+	return result, nil
+}
+
+func (s *Service) getDashboards(ctx context.Context) ([]dashboards.Dashboard, error) {
+	dashs, err := s.dashboardService.GetAllDashboards(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	result := make([]dashboards.Dashboard, len(dashs))
+	for i, dashboard := range dashs {
+		result[i] = *dashboard
+	}
+
+	return result, nil
+}
+
+// asynchronous process for writing the snapshot to the filesystem and updating the snapshot status
+func (s *Service) buildSnapshot(ctx context.Context, snapshotMeta cloudmigration.CloudMigrationSnapshot) {
+	// TODO -- make sure we can only build one snapshot at a time
+	s.buildSnapshotMutex.Lock()
+	defer s.buildSnapshotMutex.Unlock()
+	s.buildSnapshotError = false
+
+	// update snapshot status to creating, add some retries since this is a background task
+	if err := retryer.Retry(func() (retryer.RetrySignal, error) {
+		err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:    snapshotMeta.UID,
+			Status: cloudmigration.SnapshotStatusCreating,
+		})
+		return retryer.FuncComplete, err
+	}, 10, time.Millisecond*100, time.Second*10); err != nil {
+		s.log.Error("failed to set snapshot status to 'creating'", "err", err)
+		s.buildSnapshotError = true
+		return
+	}
+
+	// build snapshot
+	// just sleep for now to simulate snapshot creation happening
+	// need to do a couple of fancy things when we implement this:
+	//   - some sort of regular check-in so we know we haven't timed out
+	//   - a channel to listen for cancel events
+	//   - retries baked into the snapshot writing process?
+	s.log.Debug("snapshot meta", "snapshot", snapshotMeta)
+	time.Sleep(3 * time.Second)
+
+	// update snapshot status to pending upload with retry
+	if err := retryer.Retry(func() (retryer.RetrySignal, error) {
+		err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:    snapshotMeta.UID,
+			Status: cloudmigration.SnapshotStatusPendingUpload,
+		})
+		return retryer.FuncComplete, err
+	}, 10, time.Millisecond*100, time.Second*10); err != nil {
+		s.log.Error("failed to set snapshot status to 'pending upload'", "err", err)
+		s.buildSnapshotError = true
+	}
+}
+
+// asynchronous process for and updating the snapshot status
+func (s *Service) uploadSnapshot(ctx context.Context, snapshotMeta cloudmigration.CloudMigrationSnapshot) {
+	// TODO -- make sure we can only upload one snapshot at a time
+	s.buildSnapshotMutex.Lock()
+	defer s.buildSnapshotMutex.Unlock()
+	s.buildSnapshotError = false
+
+	// update snapshot status to uploading, add some retries since this is a background task
+	if err := retryer.Retry(func() (retryer.RetrySignal, error) {
+		err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:    snapshotMeta.UID,
+			Status: cloudmigration.SnapshotStatusUploading,
+		})
+		return retryer.FuncComplete, err
+	}, 10, time.Millisecond*100, time.Second*10); err != nil {
+		s.log.Error("failed to set snapshot status to 'creating'", "err", err)
+		s.buildSnapshotError = true
+		return
+	}
+
+	// upload snapshot
+	// just sleep for now to simulate snapshot creation happening
+	s.log.Debug("snapshot meta", "snapshot", snapshotMeta)
+	time.Sleep(3 * time.Second)
+
+	// update snapshot status to pending processing with retry
+	if err := retryer.Retry(func() (retryer.RetrySignal, error) {
+		err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:    snapshotMeta.UID,
+			Status: cloudmigration.SnapshotStatusPendingProcessing,
+		})
+		return retryer.FuncComplete, err
+	}, 10, time.Millisecond*100, time.Second*10); err != nil {
+		s.log.Error("failed to set snapshot status to 'pending upload'", "err", err)
+		s.buildSnapshotError = true
+	}
+
+	// simulate the rest
+	// processing
+	time.Sleep(3 * time.Second)
+	if err := s.store.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+		UID:    snapshotMeta.UID,
+		Status: cloudmigration.SnapshotStatusProcessing,
+	}); err != nil {
+		s.log.Error("updating snapshot", "err", err)
+	}
+	// end here as the GetSnapshot handler will fill in the rest when called
+}

pkg/services/cloudmigration/cloudmigrationimpl/store.go

@@ -15,4 +15,9 @@ type store interface {
 	CreateMigrationRun(ctx context.Context, cmr cloudmigration.CloudMigrationSnapshot) (string, error)
 	GetMigrationStatus(ctx context.Context, cmrUID string) (*cloudmigration.CloudMigrationSnapshot, error)
 	GetMigrationStatusList(ctx context.Context, migrationUID string) ([]*cloudmigration.CloudMigrationSnapshot, error)
+
+	CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) (string, error)
+	UpdateSnapshot(ctx context.Context, snapshot cloudmigration.UpdateSnapshotCmd) error
+	GetSnapshotByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSnapshot, error)
+	GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error)
 }

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -146,6 +146,106 @@ func (ss *sqlStore) GetMigrationStatusList(ctx context.Context, migrationUID str
 	return runs, nil
 }
 
+func (ss *sqlStore) CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) (string, error) {
+	if err := ss.encryptKey(ctx, &snapshot); err != nil {
+		return "", err
+	}
+	if snapshot.Result == nil {
+		snapshot.Result = make([]byte, 0)
+	}
+	if snapshot.UID == "" {
+		snapshot.UID = util.GenerateShortUID()
+	}
+
+	err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		snapshot.Created = time.Now()
+		snapshot.Updated = time.Now()
+		snapshot.UID = util.GenerateShortUID()
+
+		_, err := sess.Insert(&snapshot)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return "", err
+	}
+	return snapshot.UID, nil
+}
+
+// UpdateSnapshot takes a snapshot object containing a uid and updates a subset of features in the database.
+func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.UpdateSnapshotCmd) error {
+	if update.UID == "" {
+		return fmt.Errorf("missing snapshot uid")
+	}
+	err := ss.db.InTransaction(ctx, func(ctx context.Context) error {
+		// Update status if set
+		if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			if update.Status != "" {
+				rawSQL := "UPDATE cloud_migration_snapshot SET status=? WHERE uid=?"
+				if _, err := sess.Exec(rawSQL, update.Status, update.UID); err != nil {
+					return fmt.Errorf("updating snapshot status for uid %s: %w", update.UID, err)
+				}
+			}
+			return nil
+		}); err != nil {
+			return err
+		}
+
+		// Update result if set
+		if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			if len(update.Result) > 0 {
+				rawSQL := "UPDATE cloud_migration_snapshot SET result=? WHERE uid=?"
+				if _, err := sess.Exec(rawSQL, update.Result, update.UID); err != nil {
+					return fmt.Errorf("updating snapshot result for uid %s: %w", update.UID, err)
+				}
+			}
+			return nil
+		}); err != nil {
+			return err
+		}
+
+		return nil
+	})
+
+	return err
+}
+
+func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSnapshot, error) {
+	var snapshot cloudmigration.CloudMigrationSnapshot
+	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
+		exist, err := sess.Where("uid=?", uid).Get(&snapshot)
+		if err != nil {
+			return err
+		}
+		if !exist {
+			return cloudmigration.ErrSnapshotNotFound
+		}
+		return nil
+	})
+
+	if err := ss.decryptKey(ctx, &snapshot); err != nil {
+		return &snapshot, err
+	}
+
+	return &snapshot, err
+}
+
+func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
+	var runs = make([]cloudmigration.CloudMigrationSnapshot, 0)
+	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
+		sess.Limit(query.Limit, query.Offset)
+		return sess.Find(&runs, &cloudmigration.CloudMigrationSnapshot{
+			SessionUID: query.SessionUID,
+		})
+	})
+	if err != nil {
+		return nil, err
+	}
+	return runs, 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 {
@@ -171,3 +271,29 @@ func (ss *sqlStore) decryptToken(ctx context.Context, cm *cloudmigration.CloudMi
 
 	return nil
 }
+
+func (ss *sqlStore) encryptKey(ctx context.Context, snapshot *cloudmigration.CloudMigrationSnapshot) error {
+	s, err := ss.secretsService.Encrypt(ctx, []byte(snapshot.EncryptionKey), secrets.WithoutScope())
+	if err != nil {
+		return fmt.Errorf("encrypting key: %w", err)
+	}
+
+	snapshot.EncryptionKey = base64.StdEncoding.EncodeToString(s)
+
+	return nil
+}
+
+func (ss *sqlStore) decryptKey(ctx context.Context, snapshot *cloudmigration.CloudMigrationSnapshot) error {
+	decoded, err := base64.StdEncoding.DecodeString(snapshot.EncryptionKey)
+	if err != nil {
+		return fmt.Errorf("key could not be decoded")
+	}
+
+	t, err := ss.secretsService.Decrypt(ctx, decoded)
+	if err != nil {
+		return fmt.Errorf("decrypting key: %w", err)
+	}
+	snapshot.EncryptionKey = string(t)
+
+	return nil
+}

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store_test.go

@@ -152,7 +152,6 @@ func Test_GetMigrationStatusList(t *testing.T) {
 		list, err := s.GetMigrationStatusList(ctx, "qwerty")
 		require.NoError(t, err)
 		require.Equal(t, 2, len(list))
-		// TODO validate that this is ok
 	})
 
 	t.Run("returns no error if migration was not found, just empty list", func(t *testing.T) {
@@ -188,11 +187,11 @@ func setUpTest(t *testing.T) (*sqlstore.SQLStore, *sqlStore) {
 	// insert cloud migration run test data
 	_, err = testDB.GetSqlxSession().Exec(ctx, `
  			INSERT INTO
- 			    cloud_migration_snapshot (session_uid, uid, result, created, updated, finished)
+ 			    cloud_migration_snapshot (session_uid, uid, result, created, updated, finished, status)
  			VALUES
- 			    ('qwerty', 'poiuy', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000'),
-  				('qwerty', 'lkjhg', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000'),
-  				('zxcvbn', 'mnbvvc', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000');
+ 			    ('qwerty', 'poiuy', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000', "finished"),
+  				('qwerty', 'lkjhg', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000', "finished"),
+  				('zxcvbn', 'mnbvvc', ?, '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000', "finished");
  		`,
 		[]byte("ERROR"),
 		[]byte("OK"),