MigrationAssistant: Restrict dashboards, folders and datasources by the org id of the signed in user (#96339)

apply security patch: main/206-202410241510.patch

pkg/services/cloudmigration/api/api.go

@@ -175,7 +175,7 @@ func (cma *CloudMigrationAPI) GetSessionList(c *contextmodel.ReqContext) respons
 	ctx, span := cma.tracer.Start(c.Req.Context(), "MigrationAPI.GetSessionList")
 	defer span.End()
 
-	sl, err := cma.cloudMigrationService.GetSessionList(ctx)
+	sl, err := cma.cloudMigrationService.GetSessionList(ctx, c.OrgID)
 	if err != nil {
 		span.SetStatus(codes.Error, "session list error")
 		span.RecordError(err)
@@ -208,7 +208,7 @@ func (cma *CloudMigrationAPI) GetSession(c *contextmodel.ReqContext) response.Re
 		return response.Error(http.StatusBadRequest, "invalid session uid", err)
 	}
 
-	s, err := cma.cloudMigrationService.GetSession(ctx, uid)
+	s, err := cma.cloudMigrationService.GetSession(ctx, c.OrgID, uid)
 	if err != nil {
 		span.SetStatus(codes.Error, "session not found")
 		span.RecordError(err)
@@ -247,6 +247,7 @@ func (cma *CloudMigrationAPI) CreateSession(c *contextmodel.ReqContext) response
 	}
 	s, err := cma.cloudMigrationService.CreateSession(ctx, cloudmigration.CloudMigrationSessionRequest{
 		AuthToken: cmd.AuthToken,
+		OrgID:     c.SignedInUser.OrgID,
 	})
 	if err != nil {
 		span.SetStatus(codes.Error, "session creation error")
@@ -285,7 +286,7 @@ func (cma *CloudMigrationAPI) DeleteSession(c *contextmodel.ReqContext) response
 		return response.ErrOrFallback(http.StatusBadRequest, "invalid session uid", err)
 	}
 
-	_, err := cma.cloudMigrationService.DeleteSession(ctx, uid)
+	_, err := cma.cloudMigrationService.DeleteSession(ctx, c.OrgID, uid)
 	if err != nil {
 		span.SetStatus(codes.Error, "session delete error")
 		span.RecordError(err)
@@ -365,6 +366,7 @@ func (cma *CloudMigrationAPI) GetSnapshot(c *contextmodel.ReqContext) response.R
 		SessionUID:  sessUid,
 		ResultPage:  c.QueryInt("resultPage"),
 		ResultLimit: c.QueryInt("resultLimit"),
+		OrgID:       c.SignedInUser.OrgID,
 	}
 	if q.ResultLimit == 0 {
 		q.ResultLimit = 100
@@ -448,6 +450,7 @@ func (cma *CloudMigrationAPI) GetSnapshotList(c *contextmodel.ReqContext) respon
 		Limit:      c.QueryInt("limit"),
 		Page:       c.QueryInt("page"),
 		Sort:       c.Query("sort"),
+		OrgID:      c.SignedInUser.OrgID,
 	}
 	if q.Limit == 0 {
 		q.Limit = 100
@@ -508,7 +511,7 @@ func (cma *CloudMigrationAPI) UploadSnapshot(c *contextmodel.ReqContext) respons
 		return response.ErrOrFallback(http.StatusBadRequest, "invalid snapshot uid", err)
 	}
 
-	if err := cma.cloudMigrationService.UploadSnapshot(ctx, sessUid, snapshotUid); err != nil {
+	if err := cma.cloudMigrationService.UploadSnapshot(ctx, c.OrgID, sessUid, snapshotUid); err != nil {
 		span.SetStatus(codes.Error, "error uploading snapshot")
 		span.RecordError(err)
 

pkg/services/cloudmigration/cloudmigration.go

@@ -17,13 +17,13 @@ type Service interface {
 	DeleteToken(ctx context.Context, uid string) error
 
 	CreateSession(ctx context.Context, req CloudMigrationSessionRequest) (*CloudMigrationSessionResponse, error)
-	GetSession(ctx context.Context, migUID string) (*CloudMigrationSession, error)
+	GetSession(ctx context.Context, orgID int64, migUID string) (*CloudMigrationSession, error)
-	DeleteSession(ctx context.Context, migUID string) (*CloudMigrationSession, error)
+	DeleteSession(ctx context.Context, orgID int64, migUID string) (*CloudMigrationSession, error)
-	GetSessionList(context.Context) (*CloudMigrationSessionListResponse, error)
+	GetSessionList(ctx context.Context, orgID int64) (*CloudMigrationSessionListResponse, error)
 
 	CreateSnapshot(ctx context.Context, signedInUser *user.SignedInUser, sessionUid string) (*CloudMigrationSnapshot, error)
 	GetSnapshot(ctx context.Context, query GetSnapshotsQuery) (*CloudMigrationSnapshot, error)
 	GetSnapshotList(ctx context.Context, query ListSnapshotsQuery) ([]CloudMigrationSnapshot, error)
-	UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error
+	UploadSnapshot(ctx context.Context, orgID int64, sessionUid string, snapshotUid string) error
 	CancelSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error
 }

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration.go

@@ -362,10 +362,10 @@ func (s *Service) DeleteToken(ctx context.Context, tokenID string) error {
 	return nil
 }
 
-func (s *Service) GetSession(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, error) {
+func (s *Service) GetSession(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error) {
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetSession")
 	defer span.End()
-	migration, err := s.store.GetMigrationSessionByUID(ctx, uid)
+	migration, err := s.store.GetMigrationSessionByUID(ctx, orgID, uid)
 	if err != nil {
 		return nil, err
 	}
@@ -373,11 +373,11 @@ func (s *Service) GetSession(ctx context.Context, uid string) (*cloudmigration.C
 	return migration, nil
 }
 
-func (s *Service) GetSessionList(ctx context.Context) (*cloudmigration.CloudMigrationSessionListResponse, error) {
+func (s *Service) GetSessionList(ctx context.Context, orgID int64) (*cloudmigration.CloudMigrationSessionListResponse, error) {
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetSessionList")
 	defer span.End()
 
-	values, err := s.store.GetCloudMigrationSessionList(ctx)
+	values, err := s.store.GetCloudMigrationSessionList(ctx, orgID)
 	if err != nil {
 		return nil, fmt.Errorf("retrieving session list from store: %w", err)
 	}
@@ -408,7 +408,7 @@ func (s *Service) CreateSession(ctx context.Context, cmd cloudmigration.CloudMig
 		return nil, cloudmigration.ErrTokenInvalid.Errorf("token could not be decoded") // don't want to leak info here
 	}
 
-	migration := token.ToMigration()
+	migration := token.ToMigration(cmd.OrgID)
 	// validate token against GMS before saving
 	if err := s.ValidateToken(ctx, migration); err != nil {
 		return nil, err
@@ -429,11 +429,11 @@ func (s *Service) CreateSession(ctx context.Context, cmd cloudmigration.CloudMig
 	}, nil
 }
 
-func (s *Service) DeleteSession(ctx context.Context, sessionUID string) (*cloudmigration.CloudMigrationSession, error) {
+func (s *Service) DeleteSession(ctx context.Context, orgID int64, sessionUID string) (*cloudmigration.CloudMigrationSession, error) {
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.DeleteSession")
 	defer span.End()
 
-	session, snapshots, err := s.store.DeleteMigrationSessionByUID(ctx, sessionUID)
+	session, snapshots, err := s.store.DeleteMigrationSessionByUID(ctx, orgID, sessionUID)
 	if err != nil {
 		s.report(ctx, session, gmsclient.EventDisconnect, 0, err)
 		return nil, fmt.Errorf("deleting migration from db for session %v: %w", sessionUID, err)
@@ -451,7 +451,7 @@ func (s *Service) CreateSnapshot(ctx context.Context, signedInUser *user.SignedI
 	defer span.End()
 
 	// fetch session for the gms auth token
-	session, err := s.store.GetMigrationSessionByUID(ctx, sessionUid)
+	session, err := s.store.GetMigrationSessionByUID(ctx, signedInUser.GetOrgID(), sessionUid)
 	if err != nil {
 		return nil, fmt.Errorf("fetching migration session for uid %s: %w", sessionUid, err)
 	}
@@ -538,13 +538,13 @@ func (s *Service) GetSnapshot(ctx context.Context, query cloudmigration.GetSnaps
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetSnapshot")
 	defer span.End()
 
-	sessionUid, snapshotUid := query.SessionUID, query.SnapshotUID
+	orgID, sessionUid, snapshotUid := query.OrgID, query.SessionUID, query.SnapshotUID
-	snapshot, err := s.store.GetSnapshotByUID(ctx, sessionUid, snapshotUid, query.ResultPage, query.ResultLimit)
+	snapshot, err := s.store.GetSnapshotByUID(ctx, orgID, sessionUid, snapshotUid, query.ResultPage, query.ResultLimit)
 	if err != nil {
 		return nil, fmt.Errorf("fetching snapshot for uid %s: %w", snapshotUid, err)
 	}
 
-	session, err := s.store.GetMigrationSessionByUID(ctx, sessionUid)
+	session, err := s.store.GetMigrationSessionByUID(ctx, orgID, sessionUid)
 	if err != nil {
 		return nil, fmt.Errorf("fetching session for uid %s: %w", sessionUid, err)
 	}
@@ -594,7 +594,7 @@ func (s *Service) GetSnapshot(ctx context.Context, query cloudmigration.GetSnaps
 		}
 
 		// Refresh the snapshot after the update
-		snapshot, err = s.store.GetSnapshotByUID(ctx, sessionUid, snapshotUid, query.ResultPage, query.ResultLimit)
+		snapshot, err = s.store.GetSnapshotByUID(ctx, orgID, sessionUid, snapshotUid, query.ResultPage, query.ResultLimit)
 		if err != nil {
 			return nil, fmt.Errorf("fetching snapshot for uid %s: %w", snapshotUid, err)
 		}
@@ -685,7 +685,7 @@ func (s *Service) GetSnapshotList(ctx context.Context, query cloudmigration.List
 	return snapshotList, nil
 }
 
-func (s *Service) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+func (s *Service) UploadSnapshot(ctx context.Context, orgID int64, sessionUid string, snapshotUid string) error {
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.UploadSnapshot",
 		trace.WithAttributes(
 			attribute.String("sessionUid", sessionUid),
@@ -695,7 +695,7 @@ func (s *Service) UploadSnapshot(ctx context.Context, sessionUid string, snapsho
 	defer span.End()
 
 	// fetch session for the gms auth token
-	session, err := s.store.GetMigrationSessionByUID(ctx, sessionUid)
+	session, err := s.store.GetMigrationSessionByUID(ctx, orgID, sessionUid)
 	if err != nil {
 		return fmt.Errorf("fetching migration session for uid %s: %w", sessionUid, err)
 	}
@@ -703,6 +703,7 @@ func (s *Service) UploadSnapshot(ctx context.Context, sessionUid string, snapsho
 	snapshot, err := s.GetSnapshot(ctx, cloudmigration.GetSnapshotsQuery{
 		SnapshotUID: snapshotUid,
 		SessionUID:  sessionUid,
+		OrgID:       orgID,
 	})
 	if err != nil {
 		return fmt.Errorf("fetching snapshot with uid %s: %w", snapshotUid, err)

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration_noop.go

@@ -29,11 +29,11 @@ func (s *NoopServiceImpl) ValidateToken(ctx context.Context, cm cloudmigration.C
 	return cloudmigration.ErrMigrationDisabled
 }
 
-func (s *NoopServiceImpl) GetSession(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, error) {
+func (s *NoopServiceImpl) GetSession(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
-func (s *NoopServiceImpl) GetSessionList(ctx context.Context) (*cloudmigration.CloudMigrationSessionListResponse, error) {
+func (s *NoopServiceImpl) GetSessionList(ctx context.Context, orgID int64) (*cloudmigration.CloudMigrationSessionListResponse, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
@@ -41,7 +41,7 @@ func (s *NoopServiceImpl) CreateSession(ctx context.Context, cm cloudmigration.C
 	return nil, cloudmigration.ErrMigrationDisabled
 }
 
-func (s *NoopServiceImpl) DeleteSession(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, error) {
+func (s *NoopServiceImpl) DeleteSession(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
@@ -57,7 +57,7 @@ func (s *NoopServiceImpl) GetSnapshotList(ctx context.Context, query cloudmigrat
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
-func (s *NoopServiceImpl) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+func (s *NoopServiceImpl) UploadSnapshot(ctx context.Context, orgID int64, sessionUid string, snapshotUid string) error {
 	return cloudmigration.ErrFeatureDisabledError
 }
 

pkg/services/cloudmigration/cloudmigrationimpl/cloudmigration_test.go

@@ -357,22 +357,19 @@ func Test_OnlyQueriesStatusFromGMSWhenRequired(t *testing.T) {
 
 func Test_DeletedDashboardsNotMigrated(t *testing.T) {
 	s := setUpServiceTest(t, false).(*Service)
+
+	/** NOTE: this is not used at the moment since we changed the service
+
 	// modify what the mock returns for just this test case
 	dashMock := s.dashboardService.(*dashboards.FakeDashboardService)
 	dashMock.On("GetAllDashboards", mock.Anything).Return(
 		[]*dashboards.Dashboard{
-			{
+			{UID: "1", OrgID: 1, Data: simplejson.New()},
-				UID:  "1",
+			{UID: "2", OrgID: 1, Data: simplejson.New(), Deleted: time.Now()},
-				Data: simplejson.New(),
-			},
-			{
-				UID:     "2",
-				Data:    simplejson.New(),
-				Deleted: time.Now(),
-			},
 		},
 		nil,
 	)
+	*/
 
 	data, err := s.getMigrationDataJSON(context.TODO(), &user.SignedInUser{OrgID: 1})
 	assert.NoError(t, err)
@@ -555,7 +552,7 @@ func TestDeleteSession(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		t.Cleanup(cancel)
 
-		session, err := s.DeleteSession(ctx, "invalid-session-uid")
+		session, err := s.DeleteSession(ctx, 2, "invalid-session-uid")
 		require.Nil(t, session)
 		require.Error(t, err)
 	})
@@ -570,6 +567,7 @@ func TestDeleteSession(t *testing.T) {
 
 		cmd := cloudmigration.CloudMigrationSessionRequest{
 			AuthToken: createTokenResp.Token,
+			OrgID:     3,
 		}
 
 		createResp, err := s.CreateSession(ctx, cmd)
@@ -577,12 +575,12 @@ func TestDeleteSession(t *testing.T) {
 		require.NotEmpty(t, createResp.UID)
 		require.NotEmpty(t, createResp.Slug)
 
-		deletedSession, err := s.DeleteSession(ctx, createResp.UID)
+		deletedSession, err := s.DeleteSession(ctx, cmd.OrgID, createResp.UID)
 		require.NoError(t, err)
 		require.NotNil(t, deletedSession)
 		require.Equal(t, deletedSession.UID, createResp.UID)
 
-		notFoundSession, err := s.GetSession(ctx, deletedSession.UID)
+		notFoundSession, err := s.GetSession(ctx, cmd.OrgID, deletedSession.UID)
 		require.ErrorIs(t, err, cloudmigration.ErrMigrationNotFound)
 		require.Nil(t, notFoundSession)
 	})
@@ -868,6 +866,21 @@ func setUpServiceTest(t *testing.T, withDashboardMock bool) cloudmigration.Servi
 		LastApplied:               time.Now().Unix(),
 	}))
 
+	// Insert test data for dashboard test, should be removed later when we move GetAllDashboardsByOrgId() to the dashboard service
+	_, err = sqlStore.GetSqlxSession().Exec(context.Background(), `
+		INSERT INTO
+			dashboard (id, org_id, data, deleted, slug, title, created, version, updated )
+		VALUES
+			(1, 1, '{}', null, 'asdf', 'ghjk', '2024-03-27 15:30:43.000' , '1','2024-03-27 15:30:43.000' ),
+			(2, 1, '{}', '2024-03-27 15:30:43.000','qwert', 'yuio', '2024-03-27 15:30:43.000' , '2','2024-03-27 15:30:43.000'),
+			(3, 2, '{}', null, 'asdf', 'ghjk', '2024-03-27 15:30:43.000' , '1','2024-03-27 15:30:43.000' ),
+			(4, 2, '{}', '2024-03-27 15:30:43.000','qwert', 'yuio', '2024-03-27 15:30:43.000' , '2','2024-03-27 15:30:43.000');
+		`,
+	)
+	if err != nil {
+		require.NoError(t, err)
+	}
+
 	s, err := ProvideService(
 		cfg,
 		httpclient.NewProvider(),

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

@@ -57,21 +57,21 @@ func (m FakeServiceImpl) CreateSession(_ context.Context, _ cloudmigration.Cloud
 	}, nil
 }
 
-func (m FakeServiceImpl) GetSession(_ context.Context, _ string) (*cloudmigration.CloudMigrationSession, error) {
+func (m FakeServiceImpl) GetSession(_ context.Context, _ int64, _ string) (*cloudmigration.CloudMigrationSession, error) {
 	if m.ReturnError {
 		return nil, fmt.Errorf("mock error")
 	}
 	return &cloudmigration.CloudMigrationSession{UID: "fake"}, nil
 }
 
-func (m FakeServiceImpl) DeleteSession(_ context.Context, _ string) (*cloudmigration.CloudMigrationSession, error) {
+func (m FakeServiceImpl) DeleteSession(_ context.Context, _ int64, _ string) (*cloudmigration.CloudMigrationSession, error) {
 	if m.ReturnError {
 		return nil, fmt.Errorf("mock error")
 	}
 	return &cloudmigration.CloudMigrationSession{UID: "fake"}, nil
 }
 
-func (m FakeServiceImpl) GetSessionList(_ context.Context) (*cloudmigration.CloudMigrationSessionListResponse, error) {
+func (m FakeServiceImpl) GetSessionList(_ context.Context, _ int64) (*cloudmigration.CloudMigrationSessionListResponse, error) {
 	if m.ReturnError {
 		return nil, fmt.Errorf("mock error")
 	}
@@ -154,7 +154,7 @@ func (m FakeServiceImpl) GetSnapshotList(ctx context.Context, query cloudmigrati
 	return cloudSnapshots, nil
 }
 
-func (m FakeServiceImpl) UploadSnapshot(ctx context.Context, sessionUid string, snapshotUid string) error {
+func (m FakeServiceImpl) UploadSnapshot(ctx context.Context, _ int64, sessionUid string, snapshotUid string) error {
 	if m.ReturnError {
 		return fmt.Errorf("mock error")
 	}

pkg/services/cloudmigration/cloudmigrationimpl/snapshot_mgmt.go

@@ -44,7 +44,7 @@ func (s *Service) getMigrationDataJSON(ctx context.Context, signedInUser *user.S
 	defer span.End()
 
 	// Data sources
-	dataSources, err := s.getDataSourceCommands(ctx)
+	dataSources, err := s.getDataSourceCommands(ctx, signedInUser)
 	if err != nil {
 		s.log.Error("Failed to get datasources", "err", err)
 		return nil, err
@@ -208,17 +208,17 @@ func (s *Service) getMigrationDataJSON(ctx context.Context, signedInUser *user.S
 	return migrationData, nil
 }
 
-func (s *Service) getDataSourceCommands(ctx context.Context) ([]datasources.AddDataSourceCommand, error) {
+func (s *Service) getDataSourceCommands(ctx context.Context, signedInUser *user.SignedInUser) ([]datasources.AddDataSourceCommand, error) {
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.getDataSourceCommands")
 	defer span.End()
 
-	dataSources, err := s.dsService.GetAllDataSources(ctx, &datasources.GetAllDataSourcesQuery{})
+	dataSources, err := s.dsService.GetDataSources(ctx, &datasources.GetDataSourcesQuery{OrgID: signedInUser.GetOrgID()})
 	if err != nil {
 		s.log.Error("Failed to get all datasources", "err", err)
 		return nil, err
 	}
 
-	result := []datasources.AddDataSourceCommand{}
+	result := make([]datasources.AddDataSourceCommand, 0, len(dataSources))
 	for _, dataSource := range dataSources {
 		// Decrypt secure json to send raw credentials
 		decryptedData, err := s.secretsService.DecryptJsonData(ctx, dataSource.SecureJsonData)
@@ -253,7 +253,7 @@ func (s *Service) getDashboardAndFolderCommands(ctx context.Context, signedInUse
 	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.getDashboardAndFolderCommands")
 	defer span.End()
 
-	dashs, err := s.dashboardService.GetAllDashboards(ctx)
+	dashs, err := s.store.GetAllDashboardsByOrgId(ctx, signedInUser.GetOrgID())
 	if err != nil {
 		return nil, nil, err
 	}
@@ -279,20 +279,21 @@ func (s *Service) getDashboardAndFolderCommands(ctx context.Context, signedInUse
 	folders, err := s.folderService.GetFolders(ctx, folder.GetFoldersQuery{
 		UIDs:             folderUids,
 		SignedInUser:     signedInUser,
+		OrgID:            signedInUser.GetOrgID(),
 		WithFullpathUIDs: true,
 	})
 	if err != nil {
 		return nil, nil, err
 	}
 
-	folderCmds := make([]folder.CreateFolderCommand, len(folders))
+	folderCmds := make([]folder.CreateFolderCommand, 0, len(folders))
-	for i, f := range folders {
+	for _, f := range folders {
-		folderCmds[i] = folder.CreateFolderCommand{
+		folderCmds = append(folderCmds, folder.CreateFolderCommand{
 			UID:         f.UID,
 			Title:       f.Title,
 			Description: f.Description,
 			ParentUID:   f.ParentUID,
-		}
+		})
 	}
 
 	return dashboardCmds, folderCmds, nil
@@ -641,6 +642,7 @@ func (s *Service) getFolderNamesForFolderUIDs(ctx context.Context, signedInUser
 	folders, err := s.folderService.GetFolders(ctx, folder.GetFoldersQuery{
 		UIDs:             folderUIDs,
 		SignedInUser:     signedInUser,
+		OrgID:            signedInUser.GetOrgID(),
 		WithFullpathUIDs: true,
 	})
 	if err != nil {
@@ -661,15 +663,18 @@ func (s *Service) getFolderNamesForFolderUIDs(ctx context.Context, signedInUser
 // getParentNames finds the parent names for resources and returns a map of data type: {data UID : parentName}
 // for dashboards, folders and library elements - the parent is the parent folder
 func (s *Service) getParentNames(ctx context.Context, signedInUser *user.SignedInUser, dashboards []dashboards.Dashboard, folders []folder.CreateFolderCommand, libraryElements []libraryElement) (map[cloudmigration.MigrateDataType]map[string](string), error) {
-	parentNamesByType := make(map[cloudmigration.MigrateDataType]map[string](string))
+	parentNamesByType := make(map[cloudmigration.MigrateDataType]map[string]string)
 	for _, dataType := range currentMigrationTypes {
 		parentNamesByType[dataType] = make(map[string]string)
 	}
 
 	// Obtain list of unique folderUIDs
-	parentFolderUIDsSet := make(map[string]struct{}, len(dashboards)+len(folders)+len(libraryElements))
+	parentFolderUIDsSet := make(map[string]struct{})
 	for _, dashboard := range dashboards {
-		parentFolderUIDsSet[dashboard.FolderUID] = struct{}{}
+		// we dont need the root folder
+		if dashboard.FolderUID != "" {
+			parentFolderUIDsSet[dashboard.FolderUID] = struct{}{}
+		}
 	}
 	for _, f := range folders {
 		parentFolderUIDsSet[f.ParentUID] = struct{}{}

pkg/services/cloudmigration/cloudmigrationimpl/store.go

@@ -4,24 +4,29 @@ import (
 	"context"
 
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
+	"github.com/grafana/grafana/pkg/services/dashboards"
 )
 
 type store interface {
 	CreateMigrationSession(ctx context.Context, session cloudmigration.CloudMigrationSession) (*cloudmigration.CloudMigrationSession, error)
-	GetMigrationSessionByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, error)
+	GetMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error)
-	GetCloudMigrationSessionList(ctx context.Context) ([]*cloudmigration.CloudMigrationSession, error)
+	GetCloudMigrationSessionList(ctx context.Context, orgID int64) ([]*cloudmigration.CloudMigrationSession, error)
 	// DeleteMigrationSessionByUID deletes the migration session, and all the related snapshot and resources.
 	// the work is done in a transaction.
-	DeleteMigrationSessionByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, []cloudmigration.CloudMigrationSnapshot, error)
+	DeleteMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, []cloudmigration.CloudMigrationSnapshot, error)
 
 	CreateSnapshot(ctx context.Context, snapshot cloudmigration.CloudMigrationSnapshot) (string, error)
 	UpdateSnapshot(ctx context.Context, snapshot cloudmigration.UpdateSnapshotCmd) error
-	GetSnapshotByUID(ctx context.Context, sessUid, id string, resultPage int, resultLimit int) (*cloudmigration.CloudMigrationSnapshot, error)
+	GetSnapshotByUID(ctx context.Context, orgID int64, sessUid, id string, resultPage int, resultLimit int) (*cloudmigration.CloudMigrationSnapshot, error)
 	GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error)
-	DeleteSnapshot(ctx context.Context, snapshotUid string) error
 
-	CreateUpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error
+	// Deleted because were not used externally
-	GetSnapshotResources(ctx context.Context, snapshotUid string, page int, limit int) ([]cloudmigration.CloudMigrationResource, error)
+	// - DeleteSnapshot(ctx context.Context, snapshotUid string) error
-	GetSnapshotResourceStats(ctx context.Context, snapshotUid string) (*cloudmigration.SnapshotResourceStats, error)
+	// - CreateUpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error
-	DeleteSnapshotResources(ctx context.Context, snapshotUid string) error
+	// - GetSnapshotResources(ctx context.Context, snapshotUid string, page int, limit int) ([]cloudmigration.CloudMigrationResource, error)
+	// - GetSnapshotResourceStats(ctx context.Context, snapshotUid string) (*cloudmigration.SnapshotResourceStats, error)
+	// - DeleteSnapshotResources(ctx context.Context, snapshotUid string) error
+
+	// TODO move this function dashboards/databases/databases.go
+	GetAllDashboardsByOrgId(ctx context.Context, orgID int64) ([]*dashboards.Dashboard, error)
 }

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/services/cloudmigration"
+	"github.com/grafana/grafana/pkg/services/dashboards"
 	"github.com/grafana/grafana/pkg/services/secrets"
 	secretskv "github.com/grafana/grafana/pkg/services/secrets/kvstore"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
@@ -29,10 +30,10 @@ const (
 	GetSnapshotListSortingLatest = "latest"
 )
 
-func (ss *sqlStore) GetMigrationSessionByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, error) {
+func (ss *sqlStore) GetMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*cloudmigration.CloudMigrationSession, error) {
 	var cm cloudmigration.CloudMigrationSession
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
-		exist, err := sess.Where("uid=?", uid).Get(&cm)
+		exist, err := sess.Where("org_id=? AND uid=?", orgID, uid).Get(&cm)
 		if err != nil {
 			return err
 		}
@@ -74,11 +75,10 @@ func (ss *sqlStore) CreateMigrationSession(ctx context.Context, migration cloudm
 	return &migration, nil
 }
 
-func (ss *sqlStore) GetCloudMigrationSessionList(ctx context.Context) ([]*cloudmigration.CloudMigrationSession, error) {
+func (ss *sqlStore) GetCloudMigrationSessionList(ctx context.Context, orgID int64) ([]*cloudmigration.CloudMigrationSession, error) {
 	var migrations = make([]*cloudmigration.CloudMigrationSession, 0)
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
-		sess.OrderBy("created DESC")
+		return sess.Where("org_id=?", orgID).OrderBy("created DESC").Find(&migrations)
-		return sess.Find(&migrations)
 	})
 	if err != nil {
 		return nil, err
@@ -95,10 +95,10 @@ func (ss *sqlStore) GetCloudMigrationSessionList(ctx context.Context) ([]*cloudm
 	return migrations, nil
 }
 
-func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, uid string) (*cloudmigration.CloudMigrationSession, []cloudmigration.CloudMigrationSnapshot, error) {
+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 {
-		exist, err := sess.Where("uid=?", uid).Get(&c)
+		exist, err := sess.Where("org_id=? AND uid=?", orgID, uid).Get(&c)
 		if err != nil {
 			return err
 		}
@@ -124,11 +124,11 @@ func (ss *sqlStore) DeleteMigrationSessionByUID(ctx context.Context, uid string)
 
 	err = ss.db.InTransaction(ctx, func(ctx context.Context) error {
 		for _, snapshot := range snapshots {
-			err := ss.DeleteSnapshotResources(ctx, snapshot.UID)
+			err := ss.deleteSnapshotResources(ctx, snapshot.UID)
 			if err != nil {
 				return fmt.Errorf("deleting snapshot resource from db: %w", err)
 			}
-			err = ss.DeleteSnapshot(ctx, snapshot.UID)
+			err = ss.deleteSnapshot(ctx, orgID, snapshot.UID)
 			if err != nil {
 				return fmt.Errorf("deleting snapshot from db: %w", err)
 			}
@@ -214,7 +214,7 @@ 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 err := ss.createUpdateSnapshotResources(ctx, update.UID, update.Resources); err != nil {
 				return err
 			}
 		}
@@ -224,7 +224,7 @@ func (ss *sqlStore) UpdateSnapshot(ctx context.Context, update cloudmigration.Up
 	return err
 }
 
-func (ss *sqlStore) DeleteSnapshot(ctx context.Context, snapshotUid string) error {
+func (ss *sqlStore) deleteSnapshot(ctx context.Context, orgID int64, snapshotUid string) error {
 	return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		_, err := sess.Delete(cloudmigration.CloudMigrationSnapshot{
 			UID: snapshotUid,
@@ -233,9 +233,16 @@ func (ss *sqlStore) DeleteSnapshot(ctx context.Context, snapshotUid string) erro
 	})
 }
 
-func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, sessionUid, uid string, resultPage int, resultLimit int) (*cloudmigration.CloudMigrationSnapshot, error) {
+func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, orgID int64, sessionUid, uid string, resultPage int, resultLimit int) (*cloudmigration.CloudMigrationSnapshot, error) {
+	// first we check if the session exists, using orgId and sessionUid
+	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 {
+	err = ss.db.WithDbSession(ctx, func(sess *db.Session) error {
 		exist, err := sess.Where("session_uid=? AND uid=?", sessionUid, uid).Get(&snapshot)
 		if err != nil {
 			return err
@@ -257,11 +264,11 @@ func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, sessionUid, uid string
 		snapshot.EncryptionKey = []byte(secret)
 	}
 
-	resources, err := ss.GetSnapshotResources(ctx, uid, resultPage, resultLimit)
+	resources, err := ss.getSnapshotResources(ctx, uid, resultPage, resultLimit)
 	if err == nil {
 		snapshot.Resources = resources
 	}
-	stats, err := ss.GetSnapshotResourceStats(ctx, uid)
+	stats, err := ss.getSnapshotResourceStats(ctx, uid)
 	if err == nil {
 		snapshot.StatsRollup = *stats
 	}
@@ -274,7 +281,9 @@ func (ss *sqlStore) GetSnapshotByUID(ctx context.Context, sessionUid, uid string
 func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.ListSnapshotsQuery) ([]cloudmigration.CloudMigrationSnapshot, error) {
 	var snapshots = make([]cloudmigration.CloudMigrationSnapshot, 0)
 	err := ss.db.WithDbSession(ctx, func(sess *db.Session) error {
-		sess.Join("INNER", "cloud_migration_session", "cloud_migration_session.uid = cloud_migration_snapshot.session_uid")
+		sess.Join("INNER", "cloud_migration_session",
+			"cloud_migration_session.uid = cloud_migration_snapshot.session_uid AND cloud_migration_session.org_id = ?", query.OrgID,
+		)
 		if query.Limit != GetAllSnapshots {
 			offset := (query.Page - 1) * query.Limit
 			sess.Limit(query.Limit, offset)
@@ -298,7 +307,7 @@ func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.Li
 			snapshot.EncryptionKey = []byte(secret)
 		}
 
-		if stats, err := ss.GetSnapshotResourceStats(ctx, snapshot.UID); err != nil {
+		if stats, err := ss.getSnapshotResourceStats(ctx, snapshot.UID); err != nil {
 			return nil, err
 		} else {
 			snapshot.StatsRollup = *stats
@@ -310,7 +319,7 @@ func (ss *sqlStore) GetSnapshotList(ctx context.Context, query cloudmigration.Li
 
 // 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 {
+func (ss *sqlStore) createUpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) error {
 	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 {
@@ -344,7 +353,7 @@ func (ss *sqlStore) CreateUpdateSnapshotResources(ctx context.Context, snapshotU
 	})
 }
 
-func (ss *sqlStore) GetSnapshotResources(ctx context.Context, snapshotUid string, page int, limit int) ([]cloudmigration.CloudMigrationResource, error) {
+func (ss *sqlStore) getSnapshotResources(ctx context.Context, snapshotUid string, page int, limit int) ([]cloudmigration.CloudMigrationResource, error) {
 	if page < 1 {
 		page = 1
 	}
@@ -366,7 +375,7 @@ func (ss *sqlStore) GetSnapshotResources(ctx context.Context, snapshotUid string
 	return resources, nil
 }
 
-func (ss *sqlStore) GetSnapshotResourceStats(ctx context.Context, snapshotUid string) (*cloudmigration.SnapshotResourceStats, error) {
+func (ss *sqlStore) getSnapshotResourceStats(ctx context.Context, snapshotUid string) (*cloudmigration.SnapshotResourceStats, error) {
 	typeCounts := make([]struct {
 		Count int    `json:"count"`
 		Type  string `json:"type"`
@@ -413,7 +422,7 @@ func (ss *sqlStore) GetSnapshotResourceStats(ctx context.Context, snapshotUid st
 	return stats, nil
 }
 
-func (ss *sqlStore) DeleteSnapshotResources(ctx context.Context, snapshotUid string) error {
+func (ss *sqlStore) deleteSnapshotResources(ctx context.Context, snapshotUid string) error {
 	return ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
 		_, err := sess.Delete(cloudmigration.CloudMigrationResource{
 			SnapshotUID: snapshotUid,
@@ -456,3 +465,19 @@ func (ss *sqlStore) decryptToken(ctx context.Context, cm *cloudmigration.CloudMi
 
 	return nil
 }
+
+// TODO move this function dashboards/databases/databases.go
+func (ss *sqlStore) GetAllDashboardsByOrgId(ctx context.Context, orgID int64) ([]*dashboards.Dashboard, error) {
+	//ctx, span := tracer.Start(ctx, "dashboards.database.GetAllDashboardsByOrgId")
+	//defer span.End()
+
+	var dashs = make([]*dashboards.Dashboard, 0)
+	err := ss.db.WithDbSession(ctx, func(session *db.Session) error {
+		// "deleted IS NULL" is to avoid deleted dashboards
+		return session.Where("org_id = ? AND deleted IS NULL", orgID).Find(&dashs)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return dashs, nil
+}

pkg/services/cloudmigration/cloudmigrationimpl/xorm_store_test.go

@@ -25,7 +25,7 @@ func Test_GetAllCloudMigrationSessions(t *testing.T) {
 	ctx := context.Background()
 
 	t.Run("get all cloud_migration_session entries", func(t *testing.T) {
-		value, err := s.GetCloudMigrationSessionList(ctx)
+		value, err := s.GetCloudMigrationSessionList(ctx, 1)
 		require.NoError(t, err)
 		require.Equal(t, 3, len(value))
 		for _, m := range value {
@@ -54,6 +54,7 @@ func Test_CreateMigrationSession(t *testing.T) {
 		cm := cloudmigration.CloudMigrationSession{
 			AuthToken:   encodeToken("token"),
 			Slug:        "fake_stack",
+			OrgID:       3,
 			StackID:     1234,
 			RegionSlug:  "fake_slug",
 			ClusterSlug: "fake_cluster_slug",
@@ -63,7 +64,7 @@ func Test_CreateMigrationSession(t *testing.T) {
 		require.NotEmpty(t, sess.ID)
 		require.NotEmpty(t, sess.UID)
 
-		getRes, err := s.GetMigrationSessionByUID(ctx, sess.UID)
+		getRes, err := s.GetMigrationSessionByUID(ctx, 3, sess.UID)
 		require.NoError(t, err)
 		require.Equal(t, sess.ID, getRes.ID)
 		require.Equal(t, sess.UID, getRes.UID)
@@ -80,13 +81,15 @@ func Test_GetMigrationSessionByUID(t *testing.T) {
 	ctx := context.Background()
 	t.Run("find session by uid", func(t *testing.T) {
 		uid := "qwerty"
-		mig, err := s.GetMigrationSessionByUID(ctx, uid)
+		orgId := int64(1)
+		mig, err := s.GetMigrationSessionByUID(ctx, orgId, uid)
 		require.NoError(t, err)
 		require.Equal(t, uid, mig.UID)
+		require.Equal(t, orgId, mig.OrgID)
 	})
 
 	t.Run("returns error if session is not found by uid", func(t *testing.T) {
-		_, err := s.GetMigrationSessionByUID(ctx, "fake_uid_1234")
+		_, err := s.GetMigrationSessionByUID(ctx, 1, "fake_uid_1234")
 		require.ErrorIs(t, cloudmigration.ErrMigrationNotFound, err)
 	})
 }
@@ -115,7 +118,10 @@ func Test_SnapshotManagement(t *testing.T) {
 	ctx := context.Background()
 
 	t.Run("tests the snapshot lifecycle", func(t *testing.T) {
-		session, err := s.CreateMigrationSession(ctx, cloudmigration.CloudMigrationSession{})
+		session, err := s.CreateMigrationSession(ctx, cloudmigration.CloudMigrationSession{
+			OrgID:     1,
+			AuthToken: encodeToken("token"),
+		})
 		require.NoError(t, err)
 
 		// create a snapshot
@@ -129,7 +135,7 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NotEmpty(t, snapshotUid)
 
 		//retrieve it from the db
-		snapshot, err := s.GetSnapshotByUID(ctx, session.UID, snapshotUid, 0, 0)
+		snapshot, err := s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, 0, 0)
 		require.NoError(t, err)
 		require.Equal(t, cloudmigration.SnapshotStatusCreating, snapshot.Status)
 
@@ -138,22 +144,22 @@ func Test_SnapshotManagement(t *testing.T) {
 		require.NoError(t, err)
 
 		//retrieve it again
-		snapshot, err = s.GetSnapshotByUID(ctx, session.UID, snapshotUid, 0, 0)
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, 0, 0)
 		require.NoError(t, err)
 		require.Equal(t, cloudmigration.SnapshotStatusCreating, snapshot.Status)
 
 		// lists snapshots and ensures it's in there
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: session.UID, Page: 1, Limit: 100})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: session.UID, OrgID: 1, Page: 1, Limit: 100})
 		require.NoError(t, err)
 		require.Len(t, snapshots, 1)
 		require.Equal(t, *snapshot, snapshots[0])
 
 		// delete snapshot
-		err = s.DeleteSnapshot(ctx, snapshotUid)
+		err = s.deleteSnapshot(ctx, 1, snapshotUid)
 		require.NoError(t, err)
 
 		// now we expect not to find the snapshot
-		snapshot, err = s.GetSnapshotByUID(ctx, session.UID, snapshotUid, 0, 0)
+		snapshot, err = s.GetSnapshotByUID(ctx, 1, session.UID, snapshotUid, 0, 0)
 		require.ErrorIs(t, err, cloudmigration.ErrSnapshotNotFound)
 		require.Nil(t, snapshot)
 	})
@@ -165,12 +171,12 @@ func Test_SnapshotResources(t *testing.T) {
 
 	t.Run("tests CRUD of snapshot resources", func(t *testing.T) {
 		// Get the default rows from the test
-		resources, err := s.GetSnapshotResources(ctx, "poiuy", 0, 100)
+		resources, err := s.getSnapshotResources(ctx, "poiuy", 0, 100)
 		assert.NoError(t, err)
 		assert.Len(t, resources, 3)
 
 		// create a new resource and update an existing resource
-		err = s.CreateUpdateSnapshotResources(ctx, "poiuy", []cloudmigration.CloudMigrationResource{
+		err = s.createUpdateSnapshotResources(ctx, "poiuy", []cloudmigration.CloudMigrationResource{
 			{
 				Type:   cloudmigration.DatasourceDataType,
 				RefID:  "mi39fj",
@@ -184,7 +190,7 @@ func Test_SnapshotResources(t *testing.T) {
 		assert.NoError(t, err)
 
 		// Get resources again
-		resources, err = s.GetSnapshotResources(ctx, "poiuy", 0, 100)
+		resources, err = s.getSnapshotResources(ctx, "poiuy", 0, 100)
 		assert.NoError(t, err)
 		assert.Len(t, resources, 4)
 		// ensure existing resource was updated
@@ -203,7 +209,7 @@ func Test_SnapshotResources(t *testing.T) {
 		}
 
 		// check stats
-		stats, err := s.GetSnapshotResourceStats(ctx, "poiuy")
+		stats, err := s.getSnapshotResourceStats(ctx, "poiuy")
 		assert.NoError(t, err)
 		assert.Equal(t, map[cloudmigration.MigrateDataType]int{
 			cloudmigration.DatasourceDataType: 2,
@@ -217,10 +223,10 @@ func Test_SnapshotResources(t *testing.T) {
 		assert.Equal(t, 4, stats.Total)
 
 		// delete snapshot resources
-		err = s.DeleteSnapshotResources(ctx, "poiuy")
+		err = s.deleteSnapshotResources(ctx, "poiuy")
 		assert.NoError(t, err)
 		// make sure they're gone
-		resources, err = s.GetSnapshotResources(ctx, "poiuy", 0, 100)
+		resources, err = s.getSnapshotResources(ctx, "poiuy", 0, 100)
 		assert.NoError(t, err)
 		assert.Len(t, resources, 0)
 	})
@@ -233,7 +239,7 @@ func TestGetSnapshotList(t *testing.T) {
 	ctx := context.Background()
 
 	t.Run("returns list of snapshots that belong to a session", func(t *testing.T) {
-		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)
 
 		ids := make([]string, 0)
@@ -246,7 +252,7 @@ func TestGetSnapshotList(t *testing.T) {
 	})
 
 	t.Run("returns only one snapshot that belongs to a session", func(t *testing.T) {
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, Page: 1, Limit: 1})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, OrgID: 1, Page: 1, Limit: 1})
 		require.NoError(t, err)
 		assert.Len(t, snapshots, 1)
 	})
@@ -258,7 +264,7 @@ func TestGetSnapshotList(t *testing.T) {
 	})
 
 	t.Run("returns paginated snapshot that belongs to a session", func(t *testing.T) {
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, Page: 2, Limit: 1})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, OrgID: 1, Page: 2, Limit: 1})
 		require.NoError(t, err)
 
 		ids := make([]string, 0)
@@ -271,7 +277,7 @@ func TestGetSnapshotList(t *testing.T) {
 	})
 
 	t.Run("returns desc sorted list of snapshots that belong to a session", func(t *testing.T) {
-		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, Page: 1, Limit: 100, Sort: "latest"})
+		snapshots, err := s.GetSnapshotList(ctx, cloudmigration.ListSnapshotsQuery{SessionUID: sessionUID, OrgID: 1, Page: 1, Limit: 100, Sort: "latest"})
 		require.NoError(t, err)
 
 		ids := make([]string, 0)
@@ -291,7 +297,7 @@ func TestGetSnapshotList(t *testing.T) {
 
 	t.Run("if the session is deleted, snapshots can't be retrieved anymore", func(t *testing.T) {
 		// Delete the session.
-		_, _, err := s.DeleteMigrationSessionByUID(ctx, sessionUID)
+		_, _, err := s.DeleteMigrationSessionByUID(ctx, 1, sessionUID)
 		require.NoError(t, err)
 
 		// Fetch the snapshots that belong to the deleted session.
@@ -363,15 +369,17 @@ 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, auth_token, slug, stack_id, region_slug, cluster_slug, created, updated)
+			cloud_migration_session (id, uid, org_id, auth_token, slug, stack_id, region_slug, cluster_slug, created, updated)
 		VALUES
-			(1,'qwerty', ?, '11111', 11111, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
+			(1,'qwerty', 1, ?, '11111', 11111, 'test', 'test', '2024-03-25 15:30:36.000', '2024-03-27 15:30:43.000'),
-			(2,'asdfgh', ?, '22222', 22222, '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', ?, '33333', 33333, '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');
  		`,
 		encodeToken("12345"),
 		encodeToken("6789"),
 		encodeToken("777"),
+		encodeToken("0987"),
 	)
 	require.NoError(t, err)
 
@@ -380,9 +388,10 @@ func setUpTest(t *testing.T) (*sqlstore.SQLStore, *sqlStore) {
 		INSERT INTO
 			cloud_migration_snapshot (session_uid, uid, 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', "finished"),
+			('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', '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");
 		`,
 	)
 	require.NoError(t, err)
@@ -400,7 +409,8 @@ func setUpTest(t *testing.T) (*sqlstore.SQLStore, *sqlStore) {
 			('mnbvde', 'poiuy', 'DATASOURCE', 'jf38gh', 'OK', ''),
 			('qwerty', 'poiuy', 'DASHBOARD', 'ejcx4d', 'ERROR', 'fake error'),
 			('zxcvbn', 'poiuy', 'FOLDER', 'fi39fj', 'PENDING', ''),
-			('4fi9sd', '39fi39', 'FOLDER', 'fi39fj', 'OK', '');
+			('4fi9sd', '39fi39', 'FOLDER', 'fi39fj', 'OK', ''),
+			('4fi9ee', 'mnbvvc_org2', 'DATASOURCE', 'fi39asd', 'OK', '');
 		`,
 	)
 	require.NoError(t, err)

pkg/services/cloudmigration/model.go

@@ -21,6 +21,7 @@ var (
 // CloudMigrationSession represents a configured migration token
 type CloudMigrationSession struct {
 	ID          int64  `xorm:"pk autoincr 'id'"`
+	OrgID       int64  `xorm:"org_id"`
 	UID         string `xorm:"uid"`
 	AuthToken   string
 	Slug        string
@@ -144,6 +145,8 @@ type CloudMigrationRunList struct {
 
 type CloudMigrationSessionRequest struct {
 	AuthToken string
+	// OrgId in the on prem instance
+	OrgID int64
 }
 
 type CloudMigrationSessionResponse struct {
@@ -159,6 +162,7 @@ type CloudMigrationSessionListResponse struct {
 
 type GetSnapshotsQuery struct {
 	SnapshotUID string
+	OrgID       int64
 	SessionUID  string
 	ResultPage  int
 	ResultLimit int
@@ -166,6 +170,7 @@ type GetSnapshotsQuery struct {
 
 type ListSnapshotsQuery struct {
 	SessionUID string
+	OrgID      int64
 	Page       int
 	Limit      int
 	Sort       string
@@ -189,13 +194,14 @@ type Base64EncodedTokenPayload struct {
 	Instance Base64HGInstance
 }
 
-func (p Base64EncodedTokenPayload) ToMigration() CloudMigrationSession {
+func (p Base64EncodedTokenPayload) ToMigration(orgID int64) CloudMigrationSession {
 	return CloudMigrationSession{
 		AuthToken:   p.Token,
 		Slug:        p.Instance.Slug,
 		StackID:     p.Instance.StackID,
 		RegionSlug:  p.Instance.RegionSlug,
 		ClusterSlug: p.Instance.ClusterSlug,
+		OrgID:       orgID,
 	}
 }
 

pkg/services/sqlstore/migrations/cloud_migrations.go

@@ -66,7 +66,7 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 	}))
 
 	// --- v2 - asynchronous workflow refactor
-	sessionTable := Table{
+	migrationSessionTable := Table{
 		Name: "cloud_migration_session",
 		Columns: []*Column{
 			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
@@ -99,7 +99,7 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 		},
 	}
 
-	addTableReplaceMigrations(mg, migrationTable, sessionTable, 2, map[string]string{
+	addTableReplaceMigrations(mg, migrationTable, migrationSessionTable, 2, map[string]string{
 		"id":           "id",
 		"uid":          "uid",
 		"auth_token":   "auth_token",
@@ -171,6 +171,11 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 		Nullable: true,
 	}))
 
+	// -- Adds org_id column for for all elements - defaults to 1 (default org)
+	mg.AddMigration("add cloud_migration_session.org_id column", NewAddColumnMigration(migrationSessionTable, &Column{
+		Name: "org_id", Type: DB_BigInt, Nullable: false, Default: "1",
+	}))
+
 	mg.AddMigration("add cloud_migration_resource.error_code column", NewAddColumnMigration(migrationResourceTable, &Column{
 		Name:     "error_code",
 		Type:     DB_Text,