CloudMigrations: Bulk update local resources (#96002)

* wip

* make tests pass

* get all tests passing

* fixes

* some small cleanup

* fix test

* convert delimiter keys to struct keys

* dont execute empty sql statement

* remove printlns

* fix unit test

* a bit more cleanup

* whoops

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/grafana/grafana/pkg/infra/db"
@@ -94,6 +95,7 @@ func (ss *sqlStore) GetCloudMigrationSessionList(ctx context.Context, orgID int6
 	return migrations, nil
 }
 
+// DeleteMigrationSessionByUID deletes the migration session, and all the related snapshot and resources the work is done in a transaction.
 func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, []cloudmigration.CloudMigrationSnapshot, error) {
 	var c cloudmigration.CloudMigrationSession
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
@@ -127,7 +129,7 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, orgID int64
 			if err != nil {
 				return fmt.Errorf("deleting snapshot resource from db: %w", err)
 			}
-			err = ss.deleteSnapshot(ctx, orgID, snapshot.UID)
+			err = ss.deleteSnapshot(ctx, snapshot.UID)
 			if err != nil {
 				return fmt.Errorf("deleting snapshot from db: %w", err)
 			}
@@ -211,19 +213,26 @@ func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.Up
 			}
 		}
 
-		// Update resources if set
-		if len(update.Resources) > 0 {
-			if err := ss.createUpdateSnapshotResources(ctx, update.UID, update.Resources); err != nil {
+		// If local resources are set, it means we have to create them for the first time
+		if len(update.LocalResourcesToCreate) > 0 {
+			if err := ss.CreateSnapshotResources(ctx, update.UID, update.LocalResourcesToCreate); err != nil {
 				return err
 			}
 		}
+		// If cloud resources are set, it means we have to update our resource local state
+		if len(update.CloudResourcesToUpdate) > 0 {
+			if err := ss.UpdateSnapshotResources(ctx, update.UID, update.CloudResourcesToUpdate); err != nil {
+				return err
+			}
+		}
+
 		return nil
 	})
 
 	return err
 }
 
-func (ss *sqlStore) deleteSnapshot(ctx context.Context, orgID int64, snapshotUid string) error {
+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,
@@ -316,30 +325,90 @@ func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.Li
 	return snapshots, nil
 }
 
-// CreateUpdateSnapshotResources either updates a migration resource for a snapshot, or creates it if it does not exist
-// If the uid is not known, it uses snapshot_uid + resource_uid as a lookup
-func (ss *sqlStore) createUpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
+// CreateSnapshotResources initializes the local state of a resources belonging to a snapshot
+func (ss *sqlStore) CreateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
+	for i := 0; i < len(resources); i++ {
+		resources[i].UID = util.GenerateShortUID()
+		// ensure snapshot_uids are consistent so that we can use in conjunction with refID for lookup later
+		resources[i].SnapshotUID = snapshotUid
+	}
+
+	err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		_, err := sess.Insert(resources)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("creating resources: %w", err)
+	}
+
+	return nil
+}
+
+// UpdateSnapshotResources updates a migration resource for a snapshot, using snapshot_uid + resource_uid as a lookup
+// It does preprocessing on the results in order to minimize the sql queries executed.
+func (ss *sqlStore) UpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
+	// refIds of resources that migrated successfully in order to update in bulk
+	okIds := make([]any, 0, len(resources))
+
+	// group any failed resources by errCode and errStr
+	type errId struct {
+		errCode cloudmigration.ResourceErrorCode
+		errStr  string
+	}
+	errorIds := make(map[errId][]any)
+
+	for _, r := range resources {
+		if r.Status == cloudmigration.ItemStatusOK {
+			okIds = append(okIds, r.RefID)
+		} else if r.Status == cloudmigration.ItemStatusError {
+			key := errId{errCode: r.ErrorCode, errStr: r.Error}
+			if ids, ok := errorIds[key]; ok {
+				errorIds[key] = append(ids, r.RefID)
+			} else {
+				errorIds[key] = []any{r.RefID}
+			}
+		}
+	}
+
+	type statement struct {
+		sql  string
+		args []any
+	}
+
+	// Prepare a sql statement for all of the OK statuses
+	var okUpdateStatement *statement
+	if len(okIds) > 0 {
+		okUpdateStatement = &statement{
+			sql:  fmt.Sprintf("UPDATE cloud_migration_resource SET status=? WHERE snapshot_uid=? AND resource_uid IN (?%s)", strings.Repeat(", ?", len(okIds)-1)),
+			args: append([]any{cloudmigration.ItemStatusOK, snapshotUid}, okIds...),
+		}
+	}
+
+	// Prepare however many sql statements are necessary for the error statuses
+	errorStatements := make([]statement, 0, len(errorIds))
+	for k, ids := range errorIds {
+		errorStatements = append(errorStatements, statement{
+			sql:  fmt.Sprintf("UPDATE cloud_migration_resource SET status=?, error_code=?, error_string=? WHERE snapshot_uid=? AND resource_uid IN (?%s)", strings.Repeat(", ?", len(ids)-1)),
+			args: append([]any{cloudmigration.ItemStatusError, k.errCode, k.errStr, snapshotUid}, ids...),
+		})
+	}
+
+	// Execute the minimum number of required statements!
+
 	return ss.db.InTransaction(ctx, func(ctx context.Context) error {
-		sql := "UPDATE cloud_migration_resource SET status=?, error_string=?, error_code=? WHERE uid=? OR (snapshot_uid=? AND resource_uid=?)"
 		err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
-			for _, r := range resources {
-				// try an update first
-				result, err := sess.Exec(sql, r.Status, r.Error, r.ErrorCode, r.UID, snapshotUid, r.RefID)
-				if err != nil {
+			if okUpdateStatement != nil {
+				if _, err := sess.Exec(append([]any{okUpdateStatement.sql}, okUpdateStatement.args...)...); err != nil {
 					return err
 				}
-				// if this had no effect, assign a uid and insert instead
-				n, err := result.RowsAffected()
-				if err != nil {
+			}
+
+			for _, q := range errorStatements {
+				if _, err := sess.Exec(append([]any{q.sql}, q.args...)...); err != nil {
 					return err
-				} else if n == 0 {
-					r.UID = util.GenerateShortUID()
-					// ensure snapshot_uids are consistent so that we can use them to query when uid isn't known
-					r.SnapshotUID = snapshotUid
-					_, err := sess.Insert(r)
-					if err != nil {
-						return err
-					}
 				}
 			}
 			return nil
@@ -364,7 +433,7 @@ func (ss *sqlStore) getSnapshotResources(ctx context.Context, snapshotUid string
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
 		offset := (page - 1) * limit
 		sess.Limit(limit, offset)
-		return sess.Find(&resources, &cloudmigration.CloudMigrationResource{
+		return sess.OrderBy("id ASC").Find(&resources, &cloudmigration.CloudMigrationResource{
 			SnapshotUID: snapshotUid,
 		})
 	})