CloudMigrations: Fix issues with snapshot resource limits (#105425) (#105624)

* CloudMigrations: Fix issues with snapshot resource limits (#105425)

* fix bulk inserts

* commit progress so cursor doesn't sabotage me

* add more tests

* get everything working

* rename variable

* update comment

* regen mocks, fix k8s list method maybe

* fix bug with duplicate entries

* lint

* Snapshots: Use slices.Chunk for batching inserts

* remove extra linebreak

---------

Co-authored-by: Matheus Macabu <macabu.matheus@gmail.com>

* manually add unit tests

* make postgres integration tests happy

---------

Co-authored-by: Matheus Macabu <macabu.matheus@gmail.com>

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"slices"
 	"strings"
 	"time"
 
@@ -28,6 +29,8 @@ const (
 	secretType                   = "cloudmigration-snapshot-encryption-key"
 	GetAllSnapshots              = -1
 	GetSnapshotListSortingLatest = "latest"
+
+	maxResourceBatchSize = 1000
 )
 
 func (ss *sqlStore) GetMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error) {
@@ -191,7 +194,9 @@ func (ss *sqlStore) CreateSnapshot(ctx context.Context, snapshot cloudmigration.
 	return snapshot.UID, nil
 }
 
-// UpdateSnapshot takes a snapshot object containing a uid and updates a subset of features in the database.
+// UpdateSnapshot takes a command containing a snapshot uid and any updates to apply to the snapshot.
+// When performing multiple updates at once (e.g. updating the status and local resources), they are executed in separate transactions in order to batch insert large datasets.
+// The status is the last thing updated, as its status ultimately determines the behavior of the API.
 func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.UpdateSnapshotCmd) error {
 	if update.UID == "" {
 		return fmt.Errorf("missing snapshot uid")
@@ -199,37 +204,35 @@ func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.Up
 	if update.SessionID == "" {
 		return fmt.Errorf("missing session uid")
 	}
-	err := ss.db.InTransaction(ctx, func(ctx context.Context) error {
-		// Update status if set
-		if update.Status != "" {
-			if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
-				rawSQL := "UPDATE cloud_migration_snapshot SET status=? WHERE session_uid=? AND uid=?"
-				if _, err := sess.Exec(rawSQL, update.Status, update.SessionID, update.UID); err != nil {
-					return fmt.Errorf("updating snapshot status for uid %s: %w", update.UID, err)
-				}
-				return nil
-			}); err != nil {
-				return err
-			}
-		}
 
-		// 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 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
-			}
+	}
+
+	// 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
-	})
+	// Update the snapshot status if set
+	if update.Status != "" {
+		if err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+			rawSQL := "UPDATE cloud_migration_snapshot SET status=? WHERE session_uid=? AND uid=?"
+			if _, err := sess.Exec(rawSQL, update.Status, update.SessionID, update.UID); err != nil {
+				return fmt.Errorf("updating snapshot status for uid %s: %w", update.UID, err)
+			}
+			return nil
+		}); err != nil {
+			return err
+		}
+	}
 
-	return err
+	return nil
 }
 
 func (ss *sqlStore) deleteSnapshot(ctx context.Context, snapshotUid string) error {
@@ -326,7 +329,18 @@ func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.Li
 }
 
 // CreateSnapshotResources initializes the local state of a resources belonging to a snapshot
+// Inserting large enough datasets causes SQL errors, so we batch the inserts
 func (ss *sqlStore) CreateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
+	for chunk := range slices.Chunk(resources, maxResourceBatchSize) {
+		if err := ss.createSnapshotResources(ctx, snapshotUid, chunk); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+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
@@ -349,7 +363,18 @@ func (ss *sqlStore) CreateSnapshotResources(ctx context.Context, snapshotUid str
 
 // 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.
+// Updating large enough datasets causes SQL errors, so we batch the updates
 func (ss *sqlStore) UpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
+	for chunk := range slices.Chunk(resources, maxResourceBatchSize) {
+		if err := ss.updateSnapshotResources(ctx, snapshotUid, chunk); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+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))
 
@@ -397,7 +422,6 @@ func (ss *sqlStore) UpdateSnapshotResources(ctx context.Context, snapshotUid str
 	}
 
 	// Execute the minimum number of required statements!
-
 	return ss.db.InTransaction(ctx, func(ctx context.Context) error {
 		err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
 			if okUpdateStatement != nil {

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store_test.go

@@ -3,9 +3,11 @@ package cloudmigrationimpl
 import (
 	"context"
 	"encoding/base64"
+	"fmt"
 	"strconv"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
 	fakeSecrets "github.com/grafana/grafana/pkg/services/secrets/fakes"
 	secretskv "github.com/grafana/grafana/pkg/services/secrets/kvstore"
@@ -95,25 +97,6 @@ func Test_GetMigrationSessionByUID(t *testing.T) {
 	})
 }
 
-/** rewrite this test using the new functions
-func Test_DeleteMigrationSession(t *testing.T) {
-	_, s := setUpTest(t)
-	ctx := context.Background()
-
-	t.Run("deletes a session from the db", func(t *testing.T) {
-		uid := "qwerty"
-		session, snapshots, err := s.DeleteMigrationSessionByUID(ctx, uid)
-		require.NoError(t, err)
-		require.Equal(t, uid, session.UID)
-		require.NotNil(t, snapshots)
-
-		// now we try to find it, should return an error
-		_, err = s.GetMigrationSessionByUID(ctx, uid)
-		require.ErrorIs(t, cloudmigration.ErrMigrationNotFound, err)
-	})
-}
-*/
-
 func Test_SnapshotManagement(t *testing.T) {
 	t.Parallel()
 
@@ -166,6 +149,82 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.ErrorIs(t, err, cloudmigration.ErrSnapshotNotFound)
 		require.Nil(t, snapshot)
 	})
+
+	t.Run("tests a snapshot with a large number of resources", func(t *testing.T) {
+		session, err := s.CreateMigrationSession(ctx, cloudmigration.CloudMigrationSession{
+			OrgID:     1,
+			AuthToken: encodeToken("token"),
+		})
+		require.NoError(t, err)
+
+		// create a snapshot
+		snapshotUid, err := s.CreateSnapshot(ctx, cloudmigration.CloudMigrationSnapshot{
+			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
+		resources := make([]cloudmigration.CloudMigrationResource, numResources)
+
+		for i := 0; i < numResources; i++ {
+			resources[i] = cloudmigration.CloudMigrationResource{
+				Name:   fmt.Sprintf("Resource %d", i),
+				Type:   cloudmigration.DashboardDataType,
+				RefID:  fmt.Sprintf("refid-%d", i),
+				Status: cloudmigration.ItemStatusPending,
+			}
+		}
+
+		// Update the snapshot with the resources to create
+		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:                    snapshotUid,
+			Status:                 cloudmigration.SnapshotStatusPendingUpload,
+			SessionID:              session.UID,
+			LocalResourcesToCreate: resources,
+		})
+		require.NoError(t, err)
+
+		// Get the Snapshot and ensure it's in the right state
+		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, 1, numResources)
+		require.NoError(t, err)
+		require.Equal(t, cloudmigration.SnapshotStatusPendingUpload, snapshot.Status)
+		require.Len(t, snapshot.Resources, numResources)
+
+		for i, r := range snapshot.Resources {
+			assert.Equal(t, cloudmigration.ItemStatusPending, r.Status)
+
+			if i%2 == 0 {
+				snapshot.Resources[i].Status = cloudmigration.ItemStatusOK
+			} else {
+				snapshot.Resources[i].Status = cloudmigration.ItemStatusError
+			}
+		}
+
+		// Update the snapshot with the resources to update
+		err = s.UpdateSnapshot(ctx, cloudmigration.UpdateSnapshotCmd{
+			UID:                    snapshotUid,
+			Status:                 cloudmigration.SnapshotStatusFinished,
+			SessionID:              session.UID,
+			CloudResourcesToUpdate: snapshot.Resources,
+		})
+		require.NoError(t, err)
+
+		// Get the Snapshot and ensure it's in the right state
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, 1, numResources)
+		require.NoError(t, err)
+		require.Equal(t, cloudmigration.SnapshotStatusFinished, snapshot.Status)
+
+		for i, r := range snapshot.Resources {
+			if i%2 == 0 {
+				assert.Equal(t, cloudmigration.ItemStatusOK, r.Status)
+			} else {
+				assert.Equal(t, cloudmigration.ItemStatusError, r.Status)
+			}
+		}
+	})
 }
 
 func Test_SnapshotResources(t *testing.T) {
@@ -244,6 +303,96 @@ func Test_SnapshotResources(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Len(t, resources, 0)
 	})
+
+	t.Run("test creating and updating a large number of resources", func(t *testing.T) {
+		// Generate 50,001 test resources in order to test both update conditions (reached the batch limit or reached the end)
+		const numResources = 50001
+		resources := make([]cloudmigration.CloudMigrationResource, numResources)
+		snapshotUid := uuid.New().String()
+
+		t.Run("create the resources", func(t *testing.T) {
+			for i := 0; i < numResources; i++ {
+				resources[i] = cloudmigration.CloudMigrationResource{
+					Name:   fmt.Sprintf("Resource %d", i),
+					Type:   cloudmigration.DashboardDataType,
+					RefID:  fmt.Sprintf("refid-%d", i),
+					Status: cloudmigration.ItemStatusPending,
+				}
+			}
+
+			// Attempt to create all resources at once -- it should batch under the hood
+			err := s.CreateSnapshotResources(ctx, snapshotUid, resources)
+			require.NoError(t, err)
+
+			// Get the resources and ensure they're all there
+			resources, err := s.getSnapshotResources(ctx, snapshotUid, 1, numResources)
+			require.NoError(t, err)
+			assert.Len(t, resources, numResources)
+		})
+
+		t.Run("update the resources", func(t *testing.T) {
+			// Initially, update with a mix of ok and error statuses
+			for i := 0; i < numResources; i++ {
+				if i%2 == 0 {
+					resources[i].Status = cloudmigration.ItemStatusOK
+				} else {
+					resources[i].Status = cloudmigration.ItemStatusError
+					resources[i].ErrorCode = "test-error"
+					resources[i].Error = "test-error-message"
+				}
+			}
+
+			err := s.UpdateSnapshotResources(ctx, snapshotUid, resources)
+			require.NoError(t, err)
+
+			resources, err := s.getSnapshotResources(ctx, snapshotUid, 1, numResources)
+			require.NoError(t, err)
+			assert.Len(t, resources, numResources)
+			for i, r := range resources {
+				if i%2 == 0 {
+					assert.Equal(t, cloudmigration.ItemStatusOK, r.Status)
+				} else {
+					assert.Equal(t, cloudmigration.ItemStatusError, r.Status)
+					assert.Equal(t, "test-error", string(r.ErrorCode))
+					assert.Equal(t, "test-error-message", r.Error)
+				}
+			}
+
+			// Now update with only error statuses
+			for i := 0; i < numResources; i++ {
+				resources[i].Status = cloudmigration.ItemStatusError
+				resources[i].ErrorCode = "test-error-2"
+				resources[i].Error = "test-error-message-2"
+			}
+
+			err = s.UpdateSnapshotResources(ctx, snapshotUid, resources)
+			require.NoError(t, err)
+
+			resources, err = s.getSnapshotResources(ctx, snapshotUid, 1, numResources)
+			require.NoError(t, err)
+			assert.Len(t, resources, numResources)
+			for _, r := range resources {
+				assert.Equal(t, cloudmigration.ItemStatusError, r.Status)
+				assert.Equal(t, "test-error-2", string(r.ErrorCode))
+				assert.Equal(t, "test-error-message-2", r.Error)
+			}
+
+			// Finally, all okay
+			for i := 0; i < numResources; i++ {
+				resources[i].Status = cloudmigration.ItemStatusOK
+			}
+
+			err = s.UpdateSnapshotResources(ctx, snapshotUid, resources)
+			require.NoError(t, err)
+
+			resources, err = s.getSnapshotResources(ctx, snapshotUid, 1, numResources)
+			require.NoError(t, err)
+			assert.Len(t, resources, numResources)
+			for _, r := range resources {
+				assert.Equal(t, cloudmigration.ItemStatusOK, r.Status)
+			}
+		})
+	})
 }
 
 func TestGetSnapshotList(t *testing.T) {
@@ -386,12 +535,12 @@ func setUpTest(t *testing.T) (*sqlstore.SQLStore, *sqlStore) {
 	// insert cloud migration test data
 	_, err := testDB.GetSqlxSession().Exec(ctx, `
 		INSERT INTO
-			cloud_migration_session (id, uid, org_id, auth_token, slug, stack_id, region_slug, cluster_slug, created, updated)
+			cloud_migration_session (uid, org_id, auth_token, slug, stack_id, region_slug, cluster_slug, created, updated)
 		VALUES
-			(1,'qwerty', 1, ?, '11111', 11111, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
-			(2,'asdfgh', 1, ?, '22222', 22222, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
-			(3,'zxcvbn', 1, ?, '33333', 33333, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
-			(4,'zxcvbn_org2', 2, ?, '33333', 33333, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000');
+			('qwerty', 1, ?, '11111', 11111, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
+			('asdfgh', 1, ?, '22222', 22222, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
+			('zxcvbn', 1, ?, '33333', 33333, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
+			('zxcvbn_org2', 2, ?, '33333', 33333, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000');
  		`,
 		encodeToken("12345"),
 		encodeToken("6789"),
@@ -403,12 +552,12 @@ 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, created, updated, finished, status)
+			cloud_migration_snapshot (session_uid, uid, created, updated, 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', "finished"),
-			('qwerty', 'lkjhg', '2024-03-26 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"),
-			('zxcvbn_org2', 'mnbvvc_org2', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', '2024-03-27 15:30:43.000', "finished");
+			('qwerty', 'poiuy',  '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', 'finished'),
+			('qwerty', 'lkjhg', '2024-03-26 15:30:36.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', 'finished'),
+			('zxcvbn_org2', 'mnbvvc_org2', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000', 'finished');
 		`,
 	)
 	require.NoError(t, err)