Chore: Refactor secrets kvstore to organize testing and migrations (#54249)

* Refactor migrations and tests for secrets kvstore

* Use fake secrets store as a shortcut on tests

* Update wire

* Use global migration logger

* Fix ds proxy tests

* Fix linting issues

* Rename data source test setup function

pkg/services/datasources/service/datasource_service.go

@@ -78,8 +78,6 @@ type DataSourceRetriever interface {
 	GetDataSource(ctx context.Context, query *datasources.GetDataSourceQuery) error
 }
 
-const secretType = "datasource"
-
 // NewNameScopeResolver provides an ScopeAttributeResolver able to
 // translate a scope prefixed with "datasources:name:" into an uid based scope.
 func NewNameScopeResolver(db DataSourceRetriever) (string, accesscontrol.ScopeAttributeResolver) {
@@ -165,7 +163,7 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *datasources.AddDataSou
 				return err
 			}
 
-			return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
+			return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, kvstore.DataSourceSecretType, string(secret))
 		}
 
 		if err := s.SQLStore.AddDataSource(ctx, cmd); err != nil {
@@ -197,7 +195,7 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *datasources.AddDataSou
 func (s *Service) DeleteDataSource(ctx context.Context, cmd *datasources.DeleteDataSourceCommand) error {
 	return s.SQLStore.InTransaction(ctx, func(ctx context.Context) error {
 		cmd.UpdateSecretFn = func() error {
-			return s.SecretsStore.Del(ctx, cmd.OrgID, cmd.Name, secretType)
+			return s.SecretsStore.Del(ctx, cmd.OrgID, cmd.Name, kvstore.DataSourceSecretType)
 		}
 
 		return s.SQLStore.DeleteDataSource(ctx, cmd)
@@ -230,13 +228,13 @@ func (s *Service) UpdateDataSource(ctx context.Context, cmd *datasources.UpdateD
 				}
 
 				if query.Result.Name != cmd.Name {
-					err := s.SecretsStore.Rename(ctx, cmd.OrgId, query.Result.Name, secretType, cmd.Name)
+					err := s.SecretsStore.Rename(ctx, cmd.OrgId, query.Result.Name, kvstore.DataSourceSecretType, cmd.Name)
 					if err != nil {
 						return err
 					}
 				}
 
-				return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, secretType, string(secret))
+				return s.SecretsStore.Set(ctx, cmd.OrgId, cmd.Name, kvstore.DataSourceSecretType, string(secret))
 			}
 		}
 
@@ -299,7 +297,7 @@ func (s *Service) GetTLSConfig(ctx context.Context, ds *datasources.DataSource,
 
 func (s *Service) DecryptedValues(ctx context.Context, ds *datasources.DataSource) (map[string]string, error) {
 	decryptedValues := make(map[string]string)
-	secret, exist, err := s.SecretsStore.Get(ctx, ds.OrgId, ds.Name, secretType)
+	secret, exist, err := s.SecretsStore.Get(ctx, ds.OrgId, ds.Name, kvstore.DataSourceSecretType)
 	if err != nil {
 		return nil, err
 	}

pkg/services/datasources/service/datasource_service_test.go

@@ -15,14 +15,16 @@ import (
 
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/infra/httpclient"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/secrets"
 	"github.com/grafana/grafana/pkg/services/secrets/fakes"
-	"github.com/grafana/grafana/pkg/services/secrets/kvstore"
-	secretsManager "github.com/grafana/grafana/pkg/services/secrets/manager"
+	secretskvs "github.com/grafana/grafana/pkg/services/secrets/kvstore"
+	secretsmng "github.com/grafana/grafana/pkg/services/secrets/manager"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/setting"
 )
 
@@ -195,8 +197,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 			Type: "Kubernetes",
 		}
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		rt1, err := dsService.GetHTTPTransport(context.Background(), &ds, provider)
@@ -229,8 +232,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 		sjson := simplejson.New()
 		sjson.Set("tlsAuthWithCACert", true)
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -277,8 +281,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 		sjson := simplejson.New()
 		sjson.Set("tlsAuth", true)
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -296,7 +301,7 @@ func TestService_GetHttpTransport(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretType, string(secureJsonData))
+		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretskvs.DataSourceSecretType, string(secureJsonData))
 		require.NoError(t, err)
 
 		rt, err := dsService.GetHTTPTransport(context.Background(), &ds, provider)
@@ -322,8 +327,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 		sjson.Set("tlsAuthWithCACert", true)
 		sjson.Set("serverName", "server-name")
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -340,7 +346,7 @@ func TestService_GetHttpTransport(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretType, string(secureJsonData))
+		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretskvs.DataSourceSecretType, string(secureJsonData))
 		require.NoError(t, err)
 
 		rt, err := dsService.GetHTTPTransport(context.Background(), &ds, provider)
@@ -364,8 +370,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 		sjson := simplejson.New()
 		sjson.Set("tlsSkipVerify", true)
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -396,8 +403,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 			"httpHeaderName1": "Authorization",
 		})
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -414,7 +422,7 @@ func TestService_GetHttpTransport(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretType, string(secureJsonData))
+		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretskvs.DataSourceSecretType, string(secureJsonData))
 		require.NoError(t, err)
 
 		headers := dsService.getCustomHeaders(sjson, map[string]string{"httpHeaderValue1": "Bearer xf5yhfkpsnmgo"})
@@ -462,8 +470,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 			"timeout": 19,
 		})
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -496,8 +505,9 @@ func TestService_GetHttpTransport(t *testing.T) {
 		sjson, err := simplejson.NewJson([]byte(`{ "sigV4Auth": true }`))
 		require.NoError(t, err)
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		ds := datasources.DataSource{
@@ -532,8 +542,9 @@ func TestService_getTimeout(t *testing.T) {
 		{jsonData: simplejson.NewFromAny(map[string]interface{}{"timeout": "2"}), expectedTimeout: 2 * time.Second},
 	}
 
-	secretsStore := kvstore.SetupTestService(t)
-	secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+	sqlStore := sqlstore.InitTestDB(t)
+	secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+	secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 	dsService := ProvideService(nil, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 	for _, tc := range testCases {
@@ -552,8 +563,9 @@ func TestService_GetDecryptedValues(t *testing.T) {
 			Type: "prometheus",
 		}
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, nil, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		jsonData := map[string]string{
@@ -577,8 +589,9 @@ func TestService_GetDecryptedValues(t *testing.T) {
 			Type: "prometheus",
 		}
 
-		secretsStore := kvstore.SetupTestService(t)
-		secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
+		sqlStore := sqlstore.InitTestDB(t)
+		secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+		secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
 		dsService := ProvideService(nil, secretsService, secretsStore, nil, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService())
 
 		jsonData := map[string]string{
@@ -587,7 +600,7 @@ func TestService_GetDecryptedValues(t *testing.T) {
 		jsonString, err := json.Marshal(jsonData)
 		require.NoError(t, err)
 
-		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretType, string(jsonString))
+		err = secretsStore.Set(context.Background(), ds.OrgId, ds.Name, secretskvs.DataSourceSecretType, string(jsonString))
 		require.NoError(t, err)
 
 		values, err := dsService.DecryptedValues(context.Background(), ds)

pkg/services/datasources/service/secrets_mig.go

@@ -1,107 +0,0 @@
-package service
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/grafana/grafana/pkg/infra/kvstore"
-	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/services/datasources"
-	"github.com/grafana/grafana/pkg/services/featuremgmt"
-)
-
-const (
-	// Not set means migration has not happened
-	secretMigrationStatusKey = "secretMigrationStatus"
-	// Migration happened with disableSecretCompatibility set to false
-	compatibleSecretMigrationValue = "compatible"
-	// Migration happened with disableSecretCompatibility set to true
-	completeSecretMigrationValue = "complete"
-)
-
-type DataSourceSecretMigrationService struct {
-	dataSourcesService datasources.DataSourceService
-	kvStore            *kvstore.NamespacedKVStore
-	features           featuremgmt.FeatureToggles
-	log                log.Logger
-}
-
-func ProvideDataSourceMigrationService(
-	dataSourcesService datasources.DataSourceService,
-	kvStore kvstore.KVStore,
-	features featuremgmt.FeatureToggles,
-) *DataSourceSecretMigrationService {
-	return &DataSourceSecretMigrationService{
-		dataSourcesService: dataSourcesService,
-		kvStore:            kvstore.WithNamespace(kvStore, 0, secretType),
-		features:           features,
-		log:                log.New("secrets.migration"),
-	}
-}
-
-func (s *DataSourceSecretMigrationService) Migrate(ctx context.Context) error {
-	migrationStatus, _, err := s.kvStore.Get(ctx, secretMigrationStatusKey)
-	if err != nil {
-		return err
-	}
-	s.log.Debug(fmt.Sprint("secret migration status is ", migrationStatus))
-	// If this flag is true, delete secrets from the legacy secrets store as they are migrated
-	disableSecretsCompatibility := s.features.IsEnabled(featuremgmt.FlagDisableSecretsCompatibility)
-	// If migration hasn't happened, migrate to unified secrets and keep copy in legacy
-	// If a complete migration happened and now backwards compatibility is enabled, copy secrets back to legacy
-	needCompatibility := migrationStatus != compatibleSecretMigrationValue && !disableSecretsCompatibility
-	// If migration hasn't happened, migrate to unified secrets and delete from legacy
-	// If a compatible migration happened and now compatibility is disabled, delete secrets from legacy
-	needMigration := migrationStatus != completeSecretMigrationValue && disableSecretsCompatibility
-
-	if needCompatibility || needMigration {
-		s.log.Debug("performing secret migration", "needs migration", needMigration, "needs compatibility", needCompatibility)
-		query := &datasources.GetAllDataSourcesQuery{}
-		err := s.dataSourcesService.GetAllDataSources(ctx, query)
-		if err != nil {
-			return err
-		}
-
-		for _, ds := range query.Result {
-			secureJsonData, err := s.dataSourcesService.DecryptedValues(ctx, ds)
-			if err != nil {
-				return err
-			}
-
-			// Secrets are set by the update data source function if the SecureJsonData is set in the command
-			// Secrets are deleted by the update data source function if the disableSecretsCompatibility flag is enabled
-			err = s.dataSourcesService.UpdateDataSource(ctx, &datasources.UpdateDataSourceCommand{
-				Id:             ds.Id,
-				OrgId:          ds.OrgId,
-				Uid:            ds.Uid,
-				Name:           ds.Name,
-				JsonData:       ds.JsonData,
-				SecureJsonData: secureJsonData,
-
-				// These are needed by the SQL function due to UseBool and MustCols
-				IsDefault:       ds.IsDefault,
-				BasicAuth:       ds.BasicAuth,
-				WithCredentials: ds.WithCredentials,
-				ReadOnly:        ds.ReadOnly,
-				User:            ds.User,
-			})
-			if err != nil {
-				return err
-			}
-		}
-
-		var newMigStatus string
-		if disableSecretsCompatibility {
-			newMigStatus = completeSecretMigrationValue
-		} else {
-			newMigStatus = compatibleSecretMigrationValue
-		}
-		err = s.kvStore.Set(ctx, secretMigrationStatusKey, newMigStatus)
-		if err != nil {
-			return err
-		}
-		s.log.Debug(fmt.Sprint("set secret migration status to ", newMigStatus))
-	}
-
-	return nil
-}

pkg/services/datasources/service/secrets_mig_test.go

@@ -1,340 +0,0 @@
-package service
-
-import (
-	"context"
-	"testing"
-
-	"github.com/grafana/grafana/pkg/infra/kvstore"
-	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
-	"github.com/grafana/grafana/pkg/services/datasources"
-	"github.com/grafana/grafana/pkg/services/featuremgmt"
-	"github.com/grafana/grafana/pkg/services/secrets/fakes"
-	secretsStore "github.com/grafana/grafana/pkg/services/secrets/kvstore"
-	secretsManager "github.com/grafana/grafana/pkg/services/secrets/manager"
-	"github.com/grafana/grafana/pkg/services/sqlstore"
-	"github.com/grafana/grafana/pkg/setting"
-	"github.com/stretchr/testify/assert"
-)
-
-func SetupTestMigrationService(t *testing.T, sqlStore *sqlstore.SQLStore, kvStore kvstore.KVStore, secretsStore secretsStore.SecretsKVStore, compatibility bool) *DataSourceSecretMigrationService {
-	t.Helper()
-	cfg := &setting.Cfg{}
-	features := featuremgmt.WithFeatures()
-	if !compatibility {
-		features = featuremgmt.WithFeatures(featuremgmt.FlagDisableSecretsCompatibility, true)
-	}
-	secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
-	dsService := ProvideService(sqlStore, secretsService, secretsStore, cfg, features, acmock.New().WithDisabled(), acmock.NewMockedPermissionsService())
-	migService := ProvideDataSourceMigrationService(dsService, kvStore, features)
-	return migService
-}
-
-func TestMigrate(t *testing.T) {
-	t.Run("should migrate from legacy to unified without compatibility", func(t *testing.T) {
-		sqlStore := sqlstore.InitTestDB(t)
-		kvStore := kvstore.ProvideService(sqlStore)
-		secretsStore := secretsStore.SetupTestService(t)
-		migService := SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, false)
-
-		dataSourceName := "Test"
-		dataSourceOrg := int64(1)
-
-		// Add test data source
-		err := sqlStore.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
-			OrgId:  dataSourceOrg,
-			Name:   dataSourceName,
-			Type:   datasources.DS_MYSQL,
-			Access: datasources.DS_ACCESS_DIRECT,
-			Url:    "http://test",
-			EncryptedSecureJsonData: map[string][]byte{
-				"password": []byte("9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"),
-			},
-		})
-		assert.NoError(t, err)
-
-		// Check if the secret json data was added
-		query := &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the migration status key is empty
-		value, exist, err := kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Check that the secret is not present on the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Run the migration
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was deleted
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.Empty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, completeSecretMigrationValue, value)
-		assert.True(t, exist)
-	})
-
-	t.Run("should migrate from legacy to unified with compatibility", func(t *testing.T) {
-		sqlStore := sqlstore.InitTestDB(t)
-		kvStore := kvstore.ProvideService(sqlStore)
-		secretsStore := secretsStore.SetupTestService(t)
-		migService := SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, true)
-
-		dataSourceName := "Test"
-		dataSourceOrg := int64(1)
-
-		// Add test data source
-		err := sqlStore.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
-			OrgId:  dataSourceOrg,
-			Name:   dataSourceName,
-			Type:   datasources.DS_MYSQL,
-			Access: datasources.DS_ACCESS_DIRECT,
-			Url:    "http://test",
-			EncryptedSecureJsonData: map[string][]byte{
-				"password": []byte("9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"),
-			},
-		})
-		assert.NoError(t, err)
-
-		// Check if the secret json data was added
-		query := &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the migration status key is empty
-		value, exist, err := kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Check that the secret is not present on the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Run the migration
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was maintained for compatibility
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, compatibleSecretMigrationValue, value)
-		assert.True(t, exist)
-	})
-
-	t.Run("should replicate from unified to legacy for compatibility", func(t *testing.T) {
-		sqlStore := sqlstore.InitTestDB(t)
-		kvStore := kvstore.ProvideService(sqlStore)
-		secretsStore := secretsStore.SetupTestService(t)
-		migService := SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, false)
-
-		dataSourceName := "Test"
-		dataSourceOrg := int64(1)
-
-		// Add test data source
-		err := sqlStore.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
-			OrgId:  dataSourceOrg,
-			Name:   dataSourceName,
-			Type:   datasources.DS_MYSQL,
-			Access: datasources.DS_ACCESS_DIRECT,
-			Url:    "http://test",
-			EncryptedSecureJsonData: map[string][]byte{
-				"password": []byte("9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"),
-			},
-		})
-		assert.NoError(t, err)
-
-		// Check if the secret json data was added
-		query := &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the migration status key is empty
-		value, exist, err := kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Check that the secret is not present on the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Run the migration without compatibility
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was deleted
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.Empty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, completeSecretMigrationValue, value)
-		assert.True(t, exist)
-
-		// Run the migration with compatibility
-		migService = SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, true)
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was re-added for compatibility
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, compatibleSecretMigrationValue, value)
-		assert.True(t, exist)
-	})
-
-	t.Run("should delete from legacy to remove compatibility", func(t *testing.T) {
-		sqlStore := sqlstore.InitTestDB(t)
-		kvStore := kvstore.ProvideService(sqlStore)
-		secretsStore := secretsStore.SetupTestService(t)
-		migService := SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, true)
-
-		dataSourceName := "Test"
-		dataSourceOrg := int64(1)
-
-		// Add test data source
-		err := sqlStore.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
-			OrgId:  dataSourceOrg,
-			Name:   dataSourceName,
-			Type:   datasources.DS_MYSQL,
-			Access: datasources.DS_ACCESS_DIRECT,
-			Url:    "http://test",
-			EncryptedSecureJsonData: map[string][]byte{
-				"password": []byte("9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"),
-			},
-		})
-		assert.NoError(t, err)
-
-		// Check if the secret json data was added
-		query := &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the migration status key is empty
-		value, exist, err := kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Check that the secret is not present on the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.Empty(t, value)
-		assert.False(t, exist)
-
-		// Run the migration with compatibility
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was maintained for compatibility
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.NotEmpty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, compatibleSecretMigrationValue, value)
-		assert.True(t, exist)
-
-		// Run the migration without compatibility
-		migService = SetupTestMigrationService(t, sqlStore, kvStore, secretsStore, false)
-		err = migService.Migrate(context.Background())
-		assert.NoError(t, err)
-
-		// Check if the secure json data was deleted
-		query = &datasources.GetDataSourceQuery{OrgId: dataSourceOrg, Name: dataSourceName}
-		err = sqlStore.GetDataSource(context.Background(), query)
-		assert.NoError(t, err)
-		assert.NotNil(t, query.Result)
-		assert.Empty(t, query.Result.SecureJsonData)
-
-		// Check if the secret was added to the secret store
-		value, exist, err = secretsStore.Get(context.Background(), dataSourceOrg, dataSourceName, secretType)
-		assert.NoError(t, err)
-		assert.NotEmpty(t, value)
-		assert.True(t, exist)
-
-		// Check if the migration status key was set
-		value, exist, err = kvStore.Get(context.Background(), 0, secretType, secretMigrationStatusKey)
-		assert.NoError(t, err)
-		assert.Equal(t, completeSecretMigrationValue, value)
-		assert.True(t, exist)
-	})
-}