Secrets: Implement secrets manager plugin fallback store (#54496)

* Refactor fallback to be isolated to plugin secret store

* Check for error value on replace fallback test helper

* Move ResetPlugin from test_helpers.go to plugin.go

* Add check to GetUnwrappedStoreFromCache

* Add fallback GetAll query to WithFallbackEnabled

* Add mutex lock to WithFallbackEnabled

* Add cache to fallback store

* Fix linter issues

* Fix linter issues

* Fix linter issues

pkg/services/secrets/kvstore/cache.go

@@ -2,6 +2,7 @@ package kvstore
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -9,13 +10,15 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 )
 
+var errSecretStoreIsNotCached = errors.New("SecretsKVStore is not a CachedKVStore")
+
 type CachedKVStore struct {
 	log   log.Logger
 	cache *localcache.CacheService
 	store SecretsKVStore
 }
 
-func NewCachedKVStore(store SecretsKVStore, defaultExpiration time.Duration, cleanupInterval time.Duration) *CachedKVStore {
+func WithCache(store SecretsKVStore, defaultExpiration time.Duration, cleanupInterval time.Duration) *CachedKVStore {
 	return &CachedKVStore{
 		log:   log.New("secrets.kvstore"),
 		cache: localcache.New(defaultExpiration, cleanupInterval),
@@ -81,14 +84,9 @@ func (kv *CachedKVStore) GetAll(ctx context.Context) ([]Item, error) {
 	return kv.store.GetAll(ctx)
 }
 
-func (kv *CachedKVStore) Fallback() SecretsKVStore {
+func GetUnwrappedStoreFromCache(kv SecretsKVStore) (SecretsKVStore, error) {
-	return kv.store.Fallback()
+	if cache, ok := kv.(*CachedKVStore); ok {
-}
+		return cache.store, nil
-
+	}
-func (kv *CachedKVStore) SetFallback(store SecretsKVStore) error {
+	return nil, errSecretStoreIsNotCached
-	return kv.store.SetFallback(store)
-}
-
-func (kv *CachedKVStore) GetUnwrappedStore() SecretsKVStore {
-	return kv.store
 }

pkg/services/secrets/kvstore/kvstore.go

@@ -51,16 +51,9 @@ func ProvideService(
 			}
 		} else {
 			// as the plugin is installed, SecretsKVStoreSQL is now replaced with
-			// an instance of secretsKVStorePlugin with the sql store as a fallback
+			// an instance of SecretsKVStorePlugin with the sql store as a fallback
 			// (used for migration and in case a secret is not found).
-			store = &secretsKVStorePlugin{
+			store = NewPluginSecretsKVStore(secretsPlugin, secretsService, namespacedKVStore, features, WithCache(store, 5*time.Second, 5*time.Minute), logger)
-				secretsPlugin:                  secretsPlugin,
-				secretsService:                 secretsService,
-				log:                            logger,
-				kvstore:                        namespacedKVStore,
-				backwardsCompatibilityDisabled: features.IsEnabled(featuremgmt.FlagDisableSecretsCompatibility),
-				fallback:                       store,
-			}
 		}
 	}
 
@@ -68,7 +61,7 @@ func ProvideService(
 		logger.Debug("secrets kvstore is using the default (SQL) implementation for secrets management")
 	}
 
-	return NewCachedKVStore(store, 5*time.Second, 5*time.Minute), nil
+	return WithCache(store, 5*time.Second, 5*time.Minute), nil
 }
 
 // SecretsKVStore is an interface for k/v store.
@@ -79,8 +72,6 @@ type SecretsKVStore interface {
 	Keys(ctx context.Context, orgId int64, namespace string, typ string) ([]Key, error)
 	Rename(ctx context.Context, orgId int64, namespace string, typ string, newNamespace string) error
 	GetAll(ctx context.Context) ([]Item, error)
-	Fallback() SecretsKVStore
-	SetFallback(store SecretsKVStore) error
 }
 
 // WithType returns a kvstore wrapper with fixed orgId and type.

pkg/services/secrets/kvstore/migrations/to_plugin_mig.go

@@ -13,6 +13,8 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
+var errSecretStoreIsNotPlugin = errors.New("SecretsKVStore is not a SecretsKVStorePlugin")
+
 // MigrateToPluginService This migrator will handle migration of datasource secrets (aka Unified secrets)
 // into the plugin secrets configured
 type MigrateToPluginService struct {
@@ -46,10 +48,16 @@ func (s *MigrateToPluginService) Migrate(ctx context.Context) error {
 	if err := secretskvs.EvaluateRemoteSecretsPlugin(ctx, s.manager, s.cfg); err == nil {
 		logger.Debug("starting migration of unified secrets to the plugin")
 		// we need to get the fallback store since in this scenario the secrets store would be the plugin.
-		fallbackStore := s.secretsStore.Fallback()
+		tmpStore, err := secretskvs.GetUnwrappedStoreFromCache(s.secretsStore)
-		if fallbackStore == nil {
+		if err != nil {
-			return errors.New("unable to get fallback secret store for migration")
+			tmpStore = s.secretsStore
+			logger.Warn("secret store is not cached, this is unexpected - continuing migration anyway.")
 		}
+		pluginStore, ok := tmpStore.(*secretskvs.SecretsKVStorePlugin)
+		if !ok {
+			return errSecretStoreIsNotPlugin
+		}
+		fallbackStore := pluginStore.Fallback()
 
 		// before we start migrating, check see if plugin startup failures were already fatal
 		namespacedKVStore := secretskvs.GetNamespacedKVStore(s.kvstore)
@@ -58,22 +66,34 @@ func (s *MigrateToPluginService) Migrate(ctx context.Context) error {
 			logger.Warn("unable to determine whether plugin startup failures are fatal - continuing migration anyway.")
 		}
 
-		allSec, err := fallbackStore.GetAll(ctx)
+		var allSec []secretskvs.Item
-		if err != nil {
+		var totalSec int
-			return nil
+		// during migration we need to have fallback enabled while we move secrets to plugin
-		}
+		err = pluginStore.WithFallbackEnabled(func() error {
-		totalSec := len(allSec)
+			// get all secrets in the fallback store
-		// We just set it again as the current secret store should be the plugin secret
+			allSec, err = fallbackStore.GetAll(ctx)
-		logger.Debug(fmt.Sprintf("Total amount of secrets to migrate: %d", totalSec))
-		for i, sec := range allSec {
-			logger.Debug(fmt.Sprintf("Migrating secret %d of %d", i+1, totalSec), "current", i+1, "secretCount", totalSec)
-			err = s.secretsStore.Set(ctx, *sec.OrgId, *sec.Namespace, *sec.Type, sec.Value)
 			if err != nil {
-				return err
+				return nil
 			}
+			totalSec := len(allSec)
+			logger.Debug(fmt.Sprintf("Total amount of secrets to migrate: %d", totalSec))
+
+			// We just set it again as the current secret store should be the plugin secret
+			for i, sec := range allSec {
+				logger.Debug(fmt.Sprintf("Migrating secret %d of %d", i+1, totalSec), "current", i+1, "secretCount", totalSec)
+				err = pluginStore.Set(ctx, *sec.OrgId, *sec.Namespace, *sec.Type, sec.Value)
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		})
+		if err != nil {
+			return err
 		}
-		logger.Debug("migrated unified secrets to plugin", "number of secrets", totalSec)
+
 		// as no err was returned, when we delete all the secrets from the sql store
+		logger.Debug("migrated unified secrets to plugin", "number of secrets", totalSec)
 		for index, sec := range allSec {
 			logger.Debug(fmt.Sprintf("Cleaning secret %d of %d", index+1, totalSec), "current", index+1, "secretCount", totalSec)
 

pkg/services/secrets/kvstore/migrations/to_plugin_mig_test.go

@@ -4,9 +4,10 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"github.com/grafana/grafana/pkg/infra/kvstore"
-	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/secrets/fakes"
 	secretskvs "github.com/grafana/grafana/pkg/services/secrets/kvstore"
 	secretsManager "github.com/grafana/grafana/pkg/services/secrets/manager"
@@ -65,14 +66,14 @@ func TestFatalPluginErr_MigrationTestWithErrorDeletingUnifiedSecrets(t *testing.
 	assert.False(t, isFatal)
 }
 
-func addSecretToSqlStore(t *testing.T, sqlSecretStore *secretskvs.SecretsKVStoreSQL, ctx context.Context, orgId int64, namespace1 string, typ string, value string) {
+func addSecretToSqlStore(t *testing.T, sqlSecretStore secretskvs.SecretsKVStore, ctx context.Context, orgId int64, namespace1 string, typ string, value string) {
 	t.Helper()
 	err := sqlSecretStore.Set(ctx, orgId, namespace1, typ, value)
 	require.NoError(t, err)
 }
 
 // validates that secrets on the sql store were deleted.
-func validateSqlSecretWasDeleted(t *testing.T, sqlSecretStore *secretskvs.SecretsKVStoreSQL, ctx context.Context, orgId int64, namespace1 string, typ string) {
+func validateSqlSecretWasDeleted(t *testing.T, sqlSecretStore secretskvs.SecretsKVStore, ctx context.Context, orgId int64, namespace1 string, typ string) {
 	t.Helper()
 	res, err := sqlSecretStore.Keys(ctx, orgId, namespace1, typ)
 	require.NoError(t, err)
@@ -88,7 +89,7 @@ func validateSecretWasStoredInPlugin(t *testing.T, secretsStore secretskvs.Secre
 }
 
 // Set up services used in migration
-func setupTestMigrateToPluginService(t *testing.T) (*MigrateToPluginService, secretskvs.SecretsKVStore, *secretskvs.SecretsKVStoreSQL) {
+func setupTestMigrateToPluginService(t *testing.T) (*MigrateToPluginService, secretskvs.SecretsKVStore, secretskvs.SecretsKVStore) {
 	t.Helper()
 
 	rawCfg := `
@@ -99,7 +100,8 @@ func setupTestMigrateToPluginService(t *testing.T) (*MigrateToPluginService, sec
 	require.NoError(t, err)
 	cfg := &setting.Cfg{Raw: raw}
 	// this would be the plugin - mocked at the moment
-	secretsStoreForPlugin := secretskvs.NewFakeSecretsKVStore()
+	fallbackStore := secretskvs.WithCache(secretskvs.NewFakeSQLSecretsKVStore(t), time.Minute*5, time.Minute*5)
+	secretsStoreForPlugin := secretskvs.WithCache(secretskvs.NewFakePluginSecretsKVStore(t, featuremgmt.WithFeatures(), fallbackStore), time.Minute*5, time.Minute*5)
 
 	// this is to init the sql secret store inside the migration
 	sqlStore := sqlstore.InitTestDB(t)
@@ -114,11 +116,7 @@ func setupTestMigrateToPluginService(t *testing.T) (*MigrateToPluginService, sec
 		manager,
 	)
 
-	secretsSql := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
+	return migratorService, secretsStoreForPlugin, fallbackStore
-
-	err = secretsStoreForPlugin.SetFallback(secretsSql)
-	require.NoError(t, err)
-	return migratorService, secretsStoreForPlugin, secretsSql
 }
 
 func setupTestMigratorServiceWithDeletionError(
@@ -128,7 +126,7 @@ func setupTestMigratorServiceWithDeletionError(
 	kvstore kvstore.KVStore,
 ) *MigrateToPluginService {
 	t.Helper()
-	secretskvs.ResetPlugin()
+	t.Cleanup(secretskvs.ResetPlugin)
 	cfg := secretskvs.SetupTestConfig(t)
 	secretsService := secretsManager.SetupTestService(t, fakes.NewFakeSecretsStore())
 	manager := secretskvs.NewFakeSecretsPluginManager(t, false)
@@ -146,7 +144,7 @@ func setupTestMigratorServiceWithDeletionError(
 	err := fallback.Set(context.Background(), orgId, str, str, "bogus")
 	require.NoError(t, err)
 	fallback.DeletionError(true)
-	err = secretskv.SetFallback(fallback)
+	err = secretskvs.ReplaceFallback(t, secretskv, fallback)
 	require.NoError(t, err)
 	return migratorService
 }

pkg/services/secrets/kvstore/plugin.go

@@ -11,6 +11,7 @@ import (
 	"github.com/grafana/grafana/pkg/plugins"
 	smp "github.com/grafana/grafana/pkg/plugins/backendplugin/secretsmanagerplugin"
 	"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/setting"
 )
@@ -22,19 +23,39 @@ var (
 	errPluginNotInstalled     = errors.New("remote secret managements plugin disabled because there is no installed plugin of type `secretsmanager`")
 )
 
-// secretsKVStorePlugin provides a key/value store backed by the Grafana plugin gRPC interface
+// SecretsKVStorePlugin provides a key/value store backed by the Grafana plugin gRPC interface
-type secretsKVStorePlugin struct {
+type SecretsKVStorePlugin struct {
+	sync.Mutex
 	log                            log.Logger
 	secretsPlugin                  smp.SecretsManagerPlugin
 	secretsService                 secrets.Service
 	kvstore                        *kvstore.NamespacedKVStore
 	backwardsCompatibilityDisabled bool
-	fallback                       SecretsKVStore
+	fallbackEnabled                bool
+	fallbackStore                  SecretsKVStore
+}
+
+func NewPluginSecretsKVStore(
+	secretsPlugin smp.SecretsManagerPlugin,
+	secretsService secrets.Service,
+	kvstore *kvstore.NamespacedKVStore,
+	features featuremgmt.FeatureToggles,
+	fallback SecretsKVStore,
+	logger log.Logger,
+) *SecretsKVStorePlugin {
+	return &SecretsKVStorePlugin{
+		secretsPlugin:                  secretsPlugin,
+		secretsService:                 secretsService,
+		log:                            logger,
+		kvstore:                        kvstore,
+		backwardsCompatibilityDisabled: features.IsEnabled(featuremgmt.FlagDisableSecretsCompatibility),
+		fallbackStore:                  fallback,
+	}
 }
 
 // Get an item from the store
 // If it is the first time a secret has been retrieved and backwards compatibility is disabled, mark plugin startup errors fatal
-func (kv *secretsKVStorePlugin) Get(ctx context.Context, orgId int64, namespace string, typ string) (string, bool, error) {
+func (kv *SecretsKVStorePlugin) Get(ctx context.Context, orgId int64, namespace string, typ string) (string, bool, error) {
 	req := &smp.GetSecretRequest{
 		KeyDescriptor: &smp.Key{
 			OrgId:     orgId,
@@ -42,15 +63,20 @@ func (kv *secretsKVStorePlugin) Get(ctx context.Context, orgId int64, namespace
 			Type:      typ,
 		},
 	}
+
 	res, err := kv.secretsPlugin.GetSecret(ctx, req)
-	if err != nil {
+	if res.UserFriendlyError != "" {
-		return "", false, err
-	} else if res.UserFriendlyError != "" {
 		err = wrapUserFriendlySecretError(res.UserFriendlyError)
 	}
 
 	if res.Exists {
-		updateFatalFlag(ctx, *kv)
+		updateFatalFlag(ctx, kv)
+	}
+
+	if kv.fallbackEnabled {
+		if err != nil || res.UserFriendlyError != "" || !res.Exists {
+			res.DecryptedValue, res.Exists, err = kv.fallbackStore.Get(ctx, orgId, namespace, typ)
+		}
 	}
 
 	return res.DecryptedValue, res.Exists, err
@@ -58,7 +84,7 @@ func (kv *secretsKVStorePlugin) Get(ctx context.Context, orgId int64, namespace
 
 // Set an item in the store
 // If it is the first time a secret has been set and backwards compatibility is disabled, mark plugin startup errors fatal
-func (kv *secretsKVStorePlugin) Set(ctx context.Context, orgId int64, namespace string, typ string, value string) error {
+func (kv *SecretsKVStorePlugin) Set(ctx context.Context, orgId int64, namespace string, typ string, value string) error {
 	req := &smp.SetSecretRequest{
 		KeyDescriptor: &smp.Key{
 			OrgId:     orgId,
@@ -73,13 +99,13 @@ func (kv *secretsKVStorePlugin) Set(ctx context.Context, orgId int64, namespace
 		err = wrapUserFriendlySecretError(res.UserFriendlyError)
 	}
 
-	updateFatalFlag(ctx, *kv)
+	updateFatalFlag(ctx, kv)
 
 	return err
 }
 
 // Del deletes an item from the store.
-func (kv *secretsKVStorePlugin) Del(ctx context.Context, orgId int64, namespace string, typ string) error {
+func (kv *SecretsKVStorePlugin) Del(ctx context.Context, orgId int64, namespace string, typ string) error {
 	req := &smp.DeleteSecretRequest{
 		KeyDescriptor: &smp.Key{
 			OrgId:     orgId,
@@ -98,7 +124,7 @@ func (kv *secretsKVStorePlugin) Del(ctx context.Context, orgId int64, namespace
 
 // Keys get all keys for a given namespace. To query for all
 // organizations the constant 'kvstore.AllOrganizations' can be passed as orgId.
-func (kv *secretsKVStorePlugin) Keys(ctx context.Context, orgId int64, namespace string, typ string) ([]Key, error) {
+func (kv *SecretsKVStorePlugin) Keys(ctx context.Context, orgId int64, namespace string, typ string) ([]Key, error) {
 	req := &smp.ListSecretsRequest{
 		KeyDescriptor: &smp.Key{
 			OrgId:     orgId,
@@ -119,7 +145,7 @@ func (kv *secretsKVStorePlugin) Keys(ctx context.Context, orgId int64, namespace
 }
 
 // Rename an item in the store
-func (kv *secretsKVStorePlugin) Rename(ctx context.Context, orgId int64, namespace string, typ string, newNamespace string) error {
+func (kv *SecretsKVStorePlugin) Rename(ctx context.Context, orgId int64, namespace string, typ string, newNamespace string) error {
 	req := &smp.RenameSecretRequest{
 		KeyDescriptor: &smp.Key{
 			OrgId:     orgId,
@@ -137,7 +163,7 @@ func (kv *secretsKVStorePlugin) Rename(ctx context.Context, orgId int64, namespa
 	return err
 }
 
-func (kv *secretsKVStorePlugin) GetAll(ctx context.Context) ([]Item, error) {
+func (kv *SecretsKVStorePlugin) GetAll(ctx context.Context) ([]Item, error) {
 	req := &smp.GetAllSecretsRequest{}
 
 	res, err := kv.secretsPlugin.GetAllSecrets(ctx, req)
@@ -150,13 +176,17 @@ func (kv *secretsKVStorePlugin) GetAll(ctx context.Context) ([]Item, error) {
 	return parseItems(res.Items), err
 }
 
-func (kv *secretsKVStorePlugin) Fallback() SecretsKVStore {
+func (kv *SecretsKVStorePlugin) Fallback() SecretsKVStore {
-	return kv.fallback
+	return kv.fallbackStore
 }
 
-func (kv *secretsKVStorePlugin) SetFallback(store SecretsKVStore) error {
+func (kv *SecretsKVStorePlugin) WithFallbackEnabled(fn func() error) error {
-	kv.fallback = store
+	kv.Lock()
-	return nil
+	defer kv.Unlock()
+	kv.fallbackEnabled = true
+	err := fn()
+	kv.fallbackEnabled = false
+	return err
 }
 
 func parseKeys(keys []*smp.Key) []Key {
@@ -181,7 +211,7 @@ func parseItems(items []*smp.Item) []Item {
 	return newItems
 }
 
-func updateFatalFlag(ctx context.Context, skv secretsKVStorePlugin) {
+func updateFatalFlag(ctx context.Context, skv *SecretsKVStorePlugin) {
 	// This function makes the most sense in here because it handles all possible scenarios:
 	//   - User changed backwards compatibility flag, so we have to migrate secrets either to or from the plugin (get or set)
 	//   - Migration is on, so we migrate secrets to the plugin (set)
@@ -247,3 +277,8 @@ func StartAndReturnPlugin(mg plugins.SecretsPluginManager, ctx context.Context)
 	}
 	return mg.SecretsManager(ctx).SecretsManager, nil
 }
+
+func ResetPlugin() {
+	fatalFlagOnce = sync.Once{}
+	startupOnce = sync.Once{}
+}

pkg/services/secrets/kvstore/plugin_test.go

@@ -26,7 +26,9 @@ func TestFatalPluginErr_PluginFailsToStartWithFatalFlagNotSet(t *testing.T) {
 	require.IsType(t, &CachedKVStore{}, p.SecretsKVStore)
 
 	cachedKv, _ := p.SecretsKVStore.(*CachedKVStore)
-	assert.IsType(t, &SecretsKVStoreSQL{}, cachedKv.GetUnwrappedStore())
+	store, err := GetUnwrappedStoreFromCache(cachedKv)
+	require.NoError(t, err)
+	assert.IsType(t, &SecretsKVStoreSQL{}, store)
 }
 
 // With fatal flag not set, store a secret in the plugin while backwards compatibility is disabled

pkg/services/secrets/kvstore/sql.go

@@ -3,7 +3,6 @@ package kvstore
 import (
 	"context"
 	"encoding/base64"
-	"errors"
 	"sync"
 	"time"
 
@@ -30,10 +29,7 @@ type cachedDecrypted struct {
 	value   string
 }
 
-var (
+var b64 = base64.RawStdEncoding
-	b64                   = base64.RawStdEncoding
-	errFallbackNotAllowed = errors.New("fallback not allowed for sql secret store")
-)
 
 func NewSQLSecretsKVStore(sqlStore sqlstore.Store, secretsService secrets.Service, logger log.Logger) *SecretsKVStoreSQL {
 	return &SecretsKVStoreSQL{
@@ -244,14 +240,6 @@ func (kv *SecretsKVStoreSQL) GetAll(ctx context.Context) ([]Item, error) {
 	return items, err
 }
 
-func (kv *SecretsKVStoreSQL) Fallback() SecretsKVStore {
-	return nil
-}
-
-func (kv *SecretsKVStoreSQL) SetFallback(_ SecretsKVStore) error {
-	return errFallbackNotAllowed
-}
-
 func (kv *SecretsKVStoreSQL) getDecryptedValue(ctx context.Context, item Item) ([]byte, error) {
 	kv.decryptionCache.Lock()
 	defer kv.decryptionCache.Unlock()

pkg/services/secrets/kvstore/test_helpers.go

@@ -7,11 +7,13 @@ import (
 	"testing"
 
 	"github.com/grafana/grafana/pkg/infra/kvstore"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/plugins"
 	"github.com/grafana/grafana/pkg/plugins/backendplugin"
 	"github.com/grafana/grafana/pkg/plugins/backendplugin/secretsmanagerplugin"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/secrets/fakes"
+	secretsmng "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/require"
@@ -19,6 +21,24 @@ import (
 	"gopkg.in/ini.v1"
 )
 
+func NewFakeSQLSecretsKVStore(t *testing.T) *SecretsKVStoreSQL {
+	t.Helper()
+	sqlStore := sqlstore.InitTestDB(t)
+	secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+	return NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
+}
+
+func NewFakePluginSecretsKVStore(t *testing.T, features featuremgmt.FeatureToggles, fallback SecretsKVStore) *SecretsKVStorePlugin {
+	t.Helper()
+	sqlStore := sqlstore.InitTestDB(t)
+	secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+	store := kvstore.ProvideService(sqlStore)
+	namespacedKVStore := GetNamespacedKVStore(store)
+	manager := NewFakeSecretsPluginManager(t, false)
+	plugin := manager.SecretsManager(context.Background()).SecretsManager
+	return NewPluginSecretsKVStore(plugin, secretsService, namespacedKVStore, features, fallback, log.New("test.logger"))
+}
+
 // In memory kv store used for testing
 type FakeSecretsKVStore struct {
 	store    map[Key]string
@@ -255,9 +275,7 @@ func SetupFatalCrashTest(
 	features := NewFakeFeatureToggles(t, isBackwardsCompatDisabled)
 	manager := NewFakeSecretsPluginManager(t, shouldFailOnStart)
 	svc, err := ProvideService(sqlStore, secretService, manager, kvstore, features, cfg)
-	t.Cleanup(func() {
+	t.Cleanup(ResetPlugin)
-		fatalFlagOnce = sync.Once{}
-	})
 	return fatalCrashTestFields{
 		SecretsKVStore: svc,
 		PluginManager:  manager,
@@ -284,7 +302,14 @@ func SetupTestConfig(t *testing.T) *setting.Cfg {
 	return &setting.Cfg{Raw: raw}
 }
 
-func ResetPlugin() {
+func ReplaceFallback(t *testing.T, kv SecretsKVStore, fb SecretsKVStore) error {
-	fatalFlagOnce = sync.Once{}
+	t.Helper()
-	startupOnce = sync.Once{}
+	if store, ok := kv.(*CachedKVStore); ok {
+		kv = store.store
+	}
+	if store, ok := kv.(*SecretsKVStorePlugin); ok {
+		store.fallbackStore = fb
+		return nil
+	}
+	return errors.New("not a plugin store")
 }