CloudMigrations: Query GMS for snapshot status with a results offset (#90453)

* query GMS for status with an offset

* remove unused state

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go

@@ -524,9 +524,11 @@ func (s *Service) GetSnapshot(ctx context.Context, query cloudmigration.GetSnaps
 		return nil, fmt.Errorf("fetching session for uid %s: %w", sessionUid, err)
 	}
 
+	// Ask GMS for snapshot status while the source of truth is in the cloud
 	if snapshot.ShouldQueryGMS() {
-		// ask GMS for status if it's in the cloud
-		snapshotMeta, err := s.gmsClient.GetSnapshotStatus(ctx, *session, *snapshot)
+		// Calculate offset based on how many results we currently have responses for
+		pending := snapshot.StatsRollup.CountsByStatus[cloudmigration.ItemStatusPending]
+		snapshotMeta, err := s.gmsClient.GetSnapshotStatus(ctx, *session, *snapshot, snapshot.StatsRollup.Total-pending)
 		if err != nil {
 			return snapshot, fmt.Errorf("error fetching snapshot status from GMS: sessionUid: %s, snapshotUid: %s", sessionUid, snapshotUid)
 		}

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration_test.go

@@ -441,7 +441,7 @@ func (m *gmsClientMock) StartSnapshot(_ context.Context, _ cloudmigration.CloudM
 	return nil, nil
 }
 
-func (m *gmsClientMock) GetSnapshotStatus(_ context.Context, _ cloudmigration.CloudMigrationSession, _ cloudmigration.CloudMigrationSnapshot) (*cloudmigration.GetSnapshotStatusResponse, error) {
+func (m *gmsClientMock) GetSnapshotStatus(_ context.Context, _ cloudmigration.CloudMigrationSession, _ cloudmigration.CloudMigrationSnapshot, _ int) (*cloudmigration.GetSnapshotStatusResponse, error) {
 	m.getStatusCalled++
 	return m.getSnapshotResponse, nil
 }

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -339,7 +339,13 @@ func (ss *sqlStore) GetSnapshotResourceStats(ctx context.Context, snapshotUid st
 		Count  int    `json:"count"`
 		Status string `json:"status"`
 	}, 0)
+	total := 0
 	err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		if t, err := sess.Count(cloudmigration.CloudMigrationResource{SnapshotUID: snapshotUid}); err != nil {
+			return err
+		} else {
+			total = int(t)
+		}
 		sess.Select("count(uid) as 'count', resource_type as 'type'").
 			Table(tableName).
 			GroupBy("type").
@@ -360,6 +366,7 @@ func (ss *sqlStore) GetSnapshotResourceStats(ctx context.Context, snapshotUid st
 	stats := &cloudmigration.SnapshotResourceStats{
 		CountsByType:   make(map[cloudmigration.MigrateDataType]int, len(typeCounts)),
 		CountsByStatus: make(map[cloudmigration.ItemStatus]int, len(statusCounts)),
+		Total:          total,
 	}
 	for _, c := range typeCounts {
 		stats.CountsByType[cloudmigration.MigrateDataType(c.Type)] = c.Count

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store_test.go

@@ -257,6 +257,7 @@ func Test_SnapshotResources(t *testing.T) {
 			cloudmigration.ItemStatusOK:      3,
 			cloudmigration.ItemStatusPending: 1,
 		}, stats.CountsByStatus)
+		assert.Equal(t, 4, stats.Total)
 
 		// delete snapshot resources
 		err = s.DeleteSnapshotResources(ctx, "poiuy")