apply security patch: v11.2.x/195-202410172117.patch

commit 808c99a01896145e6627222160ce7a9c2bec4e0d
Author: Leandro Deveikis <leandro.deveikis@gmail.com>
Date:   Fri Oct 11 14:44:30 2024 -0300

    Restrict dashboards, folders and datasources by the org id of the signed in user

pkg/services/cloudmigration/api/api.go

@@ -166,7 +166,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 {
 		return response.ErrOrFallback(http.StatusInternalServerError, "session list error", err)
 	}
@@ -193,7 +193,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 {
 		return response.ErrOrFallback(http.StatusNotFound, "session not found", err)
 	}
@@ -226,6 +226,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 {
 		return response.ErrOrFallback(http.StatusInternalServerError, "session creation error", err)
@@ -260,7 +261,7 @@ func (cma *CloudMigrationAPI) RunMigration(c *contextmodel.ReqContext) response.
 		return response.ErrOrFallback(http.StatusBadRequest, "invalid migration uid", err)
 	}
 
-	result, err := cma.cloudMigrationService.RunMigration(ctx, uid)
+	result, err := cma.cloudMigrationService.RunMigration(ctx, c.OrgID, uid)
 	if err != nil {
 		return response.ErrOrFallback(http.StatusInternalServerError, "migration run error", err)
 	}
@@ -353,7 +354,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 {
 		return response.ErrOrFallback(http.StatusInternalServerError, "session delete error", err)
 	}
@@ -418,6 +419,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
@@ -491,6 +493,7 @@ func (cma *CloudMigrationAPI) GetSnapshotList(c *contextmodel.ReqContext) respon
 		SessionUID: uid,
 		Limit:      c.QueryInt("limit"),
 		Page:       c.QueryInt("page"),
+		OrgID:      c.SignedInUser.OrgID,
 	}
 	if q.Limit == 0 {
 		q.Limit = 100
@@ -542,7 +545,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 {
 		return response.ErrOrFallback(http.StatusInternalServerError, "error uploading snapshot", err)
 	}
 

pkg/services/cloudmigration/cloudmigration.go

@@ -17,17 +17,17 @@ 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)
-	DeleteSession(ctx context.Context, migUID string) (*CloudMigrationSession, error)
-	GetSessionList(context.Context) (*CloudMigrationSessionListResponse, error)
+	GetSession(ctx context.Context, orgID int64, migUID string) (*CloudMigrationSession, error)
+	DeleteSession(ctx context.Context, orgID int64, migUID string) (*CloudMigrationSession, error)
+	GetSessionList(ctx context.Context, orgID int64) (*CloudMigrationSessionListResponse, error)
 
-	RunMigration(ctx context.Context, migUID string) (*MigrateDataResponse, error)
+	RunMigration(ctx context.Context, orgID int64, migUID string) (*MigrateDataResponse, error)
 	GetMigrationStatus(ctx context.Context, runUID string) (*CloudMigrationSnapshot, error)
 	GetMigrationRunList(ctx context.Context, migUID string) (*CloudMigrationRunList, 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

@@ -337,10 +337,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) {
-	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.GetMigration")
+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
 	}
@@ -348,8 +348,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) {
-	values, err := s.store.GetCloudMigrationSessionList(ctx)
+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, orgID)
 	if err != nil {
 		return nil, fmt.Errorf("retrieving session list from store: %w", err)
 	}
@@ -380,7 +383,7 @@ func (s *Service) CreateSession(ctx context.Context, cmd cloudmigration.CloudMig
 		return nil, fmt.Errorf("invalid token") // 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, fmt.Errorf("token validation: %w", err)
@@ -401,15 +404,15 @@ func (s *Service) CreateSession(ctx context.Context, cmd cloudmigration.CloudMig
 	}, nil
 }
 
-func (s *Service) RunMigration(ctx context.Context, uid string) (*cloudmigration.MigrateDataResponse, error) {
+func (s *Service) RunMigration(ctx context.Context, orgID int64, uid string) (*cloudmigration.MigrateDataResponse, error) {
 	// Get migration to read the auth token
-	migration, err := s.GetSession(ctx, uid)
+	migration, err := s.GetSession(ctx, orgID, uid)
 	if err != nil {
 		return nil, fmt.Errorf("migration get error: %w", err)
 	}
 
 	// Get migration data JSON
-	request, err := s.getMigrationDataJSON(ctx, &user.SignedInUser{})
+	request, err := s.getMigrationDataJSON(ctx, &user.SignedInUser{OrgID: orgID})
 	if err != nil {
 		s.log.Error("error getting the json request body for migration run", "err", err.Error())
 		return nil, fmt.Errorf("migration data get error: %w", err)
@@ -469,8 +472,11 @@ func (s *Service) GetMigrationRunList(ctx context.Context, migUID string) (*clou
 	return runList, nil
 }
 
-func (s *Service) DeleteSession(ctx context.Context, sessionUID string) (*cloudmigration.CloudMigrationSession, error) {
-	session, snapshots, err := s.store.DeleteMigrationSessionByUID(ctx, sessionUID)
+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, 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)
@@ -488,7 +494,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)
 	}
@@ -565,13 +571,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
-	snapshot, err := s.store.GetSnapshotByUID(ctx, sessionUid, snapshotUid, query.ResultPage, query.ResultLimit)
+	orgID, sessionUid, snapshotUid := query.OrgID, query.SessionUID, query.SnapshotUID
+	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)
 	}
@@ -614,7 +620,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)
 		}
@@ -642,7 +648,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),
@@ -652,7 +658,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)
 	}
@@ -660,6 +666,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.ErrFeatureDisabledError
 }
 
-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
 }
 
@@ -49,7 +49,7 @@ func (s *NoopServiceImpl) GetMigrationRunList(ctx context.Context, uid string) (
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
-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) CreateMigrationRun(context.Context, cloudmigration.Clo
 	return "", cloudmigration.ErrInternalNotImplementedError
 }
 
-func (s *NoopServiceImpl) RunMigration(context.Context, string) (*cloudmigration.MigrateDataResponse, error) {
+func (s *NoopServiceImpl) RunMigration(context.Context, int64, string) (*cloudmigration.MigrateDataResponse, error) {
 	return nil, cloudmigration.ErrFeatureDisabledError
 }
 
@@ -73,7 +73,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

@@ -9,7 +9,6 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/grafana/grafana/pkg/api/routing"
-	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/kvstore"
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -31,7 +30,6 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 )
@@ -74,6 +72,7 @@ func Test_CreateGetRunMigrationsAndRuns(t *testing.T) {
 
 	cmd := cloudmigration.CloudMigrationSessionRequest{
 		AuthToken: createTokenResp.Token,
+		OrgID:     1,
 	}
 
 	createResp, err := s.CreateSession(context.Background(), cmd)
@@ -81,20 +80,20 @@ func Test_CreateGetRunMigrationsAndRuns(t *testing.T) {
 	require.NotEmpty(t, createResp.UID)
 	require.NotEmpty(t, createResp.Slug)
 
-	getMigResp, err := s.GetSession(context.Background(), createResp.UID)
+	getMigResp, err := s.GetSession(context.Background(), 1, createResp.UID)
 	require.NoError(t, err)
 	require.NotNil(t, getMigResp)
 	require.Equal(t, createResp.UID, getMigResp.UID)
 	require.Equal(t, createResp.Slug, getMigResp.Slug)
 
-	listResp, err := s.GetSessionList(context.Background())
+	listResp, err := s.GetSessionList(context.Background(), 1)
 	require.NoError(t, err)
 	require.NotNil(t, listResp)
 	require.Equal(t, 1, len(listResp.Sessions))
 	require.Equal(t, createResp.UID, listResp.Sessions[0].UID)
 	require.Equal(t, createResp.Slug, listResp.Sessions[0].Slug)
 
-	runResp, err := s.RunMigration(ctxWithSignedInUser(), createResp.UID)
+	runResp, err := s.RunMigration(ctxWithSignedInUser(), 1, createResp.UID)
 	require.NoError(t, err)
 	require.NotNil(t, runResp)
 	resultItemsByType := make(map[string]int)
@@ -375,22 +374,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",
-				Data: simplejson.New(),
-			},
-			{
-				UID:     "2",
-				Data:    simplejson.New(),
-				Deleted: time.Now(),
-			},
+			{UID: "1", OrgID: 1, Data: simplejson.New()},
+			{UID: "2", OrgID: 1, Data: simplejson.New(), Deleted: time.Now()},
 		},
 		nil,
 	)
+	*/
 
 	data, err := s.getMigrationDataJSON(context.TODO(), &user.SignedInUser{OrgID: 1})
 	assert.NoError(t, err)
@@ -555,7 +551,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 +566,7 @@ func TestDeleteSession(t *testing.T) {
 
 		cmd := cloudmigration.CloudMigrationSessionRequest{
 			AuthToken: createTokenResp.Token,
+			OrgID:     3,
 		}
 
 		createResp, err := s.CreateSession(ctx, cmd)
@@ -577,12 +574,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)
 	})
@@ -638,7 +635,7 @@ func setUpServiceTest(t *testing.T, withDashboardMock bool) cloudmigration.Servi
 	spanRecorder := tracetest.NewSpanRecorder()
 	tracer := tracing.InitializeTracerForTest(tracing.WithSpanProcessor(spanRecorder))
 	mockFolder := &foldertest.FakeService{
-		ExpectedFolder: &folder.Folder{UID: "folderUID", Title: "Folder"},
+		ExpectedFolder: &folder.Folder{UID: "folderUID", OrgID: 1, Title: "Folder"},
 	}
 
 	cfg := setting.NewCfg()
@@ -651,6 +648,7 @@ func setUpServiceTest(t *testing.T, withDashboardMock bool) cloudmigration.Servi
 	cfg.CloudMigration.SnapshotFolder = filepath.Join(os.TempDir(), uuid.NewString())
 
 	dashboardService := dashboards.NewFakeDashboardService(t)
+	/**
 	if withDashboardMock {
 		dashboardService.On("GetAllDashboards", mock.Anything).Return(
 			[]*dashboards.Dashboard{
@@ -662,14 +660,28 @@ func setUpServiceTest(t *testing.T, withDashboardMock bool) cloudmigration.Servi
 			nil,
 		)
 	}
+	*/
 
 	dsService := &datafakes.FakeDataSourceService{
 		DataSources: []*datasources.DataSource{
-			{Name: "mmm", Type: "mysql"},
-			{Name: "ZZZ", Type: "infinity"},
+			{Name: "mmm", OrgID: 1, Type: "mysql"},
+			{Name: "ZZZ", OrgID: 1, Type: "infinity"},
 		},
 	}
 
+	// Insert test data for dashboard test, should be removed later
+	_, 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');
+ 		`,
+	)
+	require.NoError(t, err)
+
 	s, err := ProvideService(
 		cfg,
 		featuremgmt.WithFeatures(

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

@@ -56,21 +56,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")
 	}
@@ -82,7 +82,7 @@ func (m FakeServiceImpl) GetSessionList(_ context.Context) (*cloudmigration.Clou
 	}, nil
 }
 
-func (m FakeServiceImpl) RunMigration(_ context.Context, _ string) (*cloudmigration.MigrateDataResponse, error) {
+func (m FakeServiceImpl) RunMigration(_ context.Context, _ int64, _ string) (*cloudmigration.MigrateDataResponse, error) {
 	if m.ReturnError {
 		return nil, fmt.Errorf("mock error")
 	}
@@ -170,7 +170,7 @@ func (m FakeServiceImpl) GetSnapshotList(ctx context.Context, query cloudmigrati
 	}, 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

@@ -25,7 +25,7 @@ import (
 
 func (s *Service) getMigrationDataJSON(ctx context.Context, signedInUser *user.SignedInUser) (*cloudmigration.MigrateDataRequest, error) {
 	// 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
@@ -85,14 +85,17 @@ func (s *Service) getMigrationDataJSON(ctx context.Context, signedInUser *user.S
 	return migrationData, nil
 }
 
-func (s *Service) getDataSourceCommands(ctx context.Context) ([]datasources.AddDataSourceCommand, error) {
-	dataSources, err := s.dsService.GetAllDataSources(ctx, &datasources.GetAllDataSourcesQuery{})
+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.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)
@@ -124,7 +127,10 @@ func (s *Service) getDataSourceCommands(ctx context.Context) ([]datasources.AddD
 
 // getDashboardAndFolderCommands returns the json payloads required by the dashboard and folder creation APIs
 func (s *Service) getDashboardAndFolderCommands(ctx context.Context, signedInUser *user.SignedInUser) ([]dashboards.Dashboard, []folder.CreateFolderCommand, error) {
-	dashs, err := s.dashboardService.GetAllDashboards(ctx)
+	ctx, span := s.tracer.Start(ctx, "CloudMigrationService.getDashboardAndFolderCommands")
+	defer span.End()
+
+	dashs, err := s.store.GetAllDashboardsByOrgId(ctx, signedInUser.GetOrgID())
 	if err != nil {
 		return nil, nil, err
 	}
@@ -150,20 +156,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))
-	for i, f := range folders {
-		folderCmds[i] = folder.CreateFolderCommand{
+	folderCmds := make([]folder.CreateFolderCommand, 0, len(folders))
+	for _, f := range folders {
+		folderCmds = append(folderCmds, folder.CreateFolderCommand{
 			UID:         f.UID,
 			Title:       f.Title,
 			Description: f.Description,
 			ParentUID:   f.ParentUID,
-		}
+		})
 	}
 
 	return dashboardCmds, folderCmds, nil

pkg/services/cloudmigration/cloudmigrationimpl/store.go

@@ -4,15 +4,16 @@ 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)
-	GetCloudMigrationSessionList(ctx context.Context) ([]*cloudmigration.CloudMigrationSession, error)
+	GetMigrationSessionByUID(ctx context.Context, orgID int64, uid string) (*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)
 
 	CreateMigrationRun(ctx context.Context, cmr cloudmigration.CloudMigrationSnapshot) (string, error)
 	GetMigrationStatus(ctx context.Context, cmrUID string) (*cloudmigration.CloudMigrationSnapshot, error)
@@ -21,12 +22,16 @@ type store interface {
 
 	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
-	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
+	// Deleted because were not used externally
+	// - DeleteSnapshot(ctx context.Context, snapshotUid string) error
+	// - CreateUpdateSnapshotResources(ctx context.Context, snapshotUid string, resources []cloudmigration.CloudMigrationResource) 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"
@@ -28,10 +29,10 @@ const (
 	GetAllSnapshots = -1
 )
 
-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
 		}
@@ -89,11 +90,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.Find(&migrations)
+		return sess.Where("org_id=?", orgID).OrderBy("created DESC").Find(&migrations)
 	})
 	if err != nil {
 		return nil, err
@@ -110,10 +110,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
 		}
@@ -139,11 +139,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)
 			}
@@ -257,7 +257,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
 			}
 		}
@@ -267,7 +267,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,
@@ -276,9 +276,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
@@ -300,11 +307,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
 	}
@@ -317,7 +324,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)
@@ -339,7 +348,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
@@ -351,7 +360,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=? WHERE uid=? OR (snapshot_uid=? AND resource_uid=?)"
 		err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
@@ -385,7 +394,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
 	}
@@ -407,7 +416,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"`
@@ -454,7 +463,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,
@@ -497,3 +506,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

@@ -26,7 +26,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 {
@@ -55,6 +55,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",
@@ -64,7 +65,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)
@@ -81,13 +82,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)
 	})
 }
@@ -171,7 +174,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
@@ -185,7 +191,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)
 
@@ -194,22 +200,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)
 	})
@@ -221,12 +227,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",
@@ -240,7 +246,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
@@ -259,7 +265,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,
@@ -273,10 +279,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)
 	})
@@ -289,7 +295,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)
@@ -310,7 +316,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.
@@ -382,15 +388,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'),
-			(2,'asdfgh', ?, '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');
+			(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');
  		`,
 		encodeToken("12345"),
 		encodeToken("6789"),
 		encodeToken("777"),
+		encodeToken("0987"),
 	)
 	require.NoError(t, err)
 
@@ -399,9 +407,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', '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");
+			('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");
 		`,
 	)
 	require.NoError(t, err)
@@ -419,7 +428,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
@@ -118,6 +119,8 @@ type CloudMigrationRunList struct {
 
 type CloudMigrationSessionRequest struct {
 	AuthToken string
+	// OrgId in the on prem instance
+	OrgID int64
 }
 
 type CloudMigrationSessionResponse struct {
@@ -133,6 +136,7 @@ type CloudMigrationSessionListResponse struct {
 
 type GetSnapshotsQuery struct {
 	SnapshotUID string
+	OrgID       int64
 	SessionUID  string
 	ResultPage  int
 	ResultLimit int
@@ -140,6 +144,7 @@ type GetSnapshotsQuery struct {
 
 type ListSnapshotsQuery struct {
 	SessionUID string
+	OrgID      int64
 	Page       int
 	Limit      int
 }
@@ -162,13 +167,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",
@@ -158,4 +158,9 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 
 	// -- delete the snapshot result column while still in the experimental phase
 	mg.AddMigration("delete cloud_migration_snapshot.result column", NewRawSQLMigration("ALTER TABLE cloud_migration_snapshot DROP COLUMN result"))
+
+	// -- 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",
+	}))
 }