feat(unified-storage): no SQL fallback for dashboards/folders in Mode 5 (#111672)

pkg/storage/unified/client.go

@@ -20,6 +20,7 @@ import (
 	"github.com/grafana/dskit/grpcclient"
 	"github.com/grafana/dskit/middleware"
 	"github.com/grafana/dskit/services"
+	grafanarest "github.com/grafana/grafana/pkg/apiserver/rest"
 	infraDB "github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	secrets "github.com/grafana/grafana/pkg/registry/apis/secret/contracts"
@@ -66,10 +67,23 @@ func ProvideUnifiedStorageClient(opts *Options,
 		BlobThresholdBytes:  apiserverCfg.Key("blob_threshold_bytes").MustInt(options.BlobThresholdDefault),
 	}, opts.Cfg, opts.Features, opts.DB, opts.Tracer, opts.Reg, opts.Authzc, opts.Docs, storageMetrics, indexMetrics, opts.SecureValues)
 	if err == nil {
+		// Decide whether to disable SQL fallback stats per resource in Mode 5.
+		// Otherwise we would still try to query the legacy SQL database in Mode 5.
+		var disableDashboardsFallback, disableFoldersFallback bool
+		if opts.Cfg != nil {
+			// String are static here, so we don't need to import the packages.
+			foldersMode := opts.Cfg.UnifiedStorage["folders.folder.grafana.app"].DualWriterMode
+			disableFoldersFallback = foldersMode == grafanarest.Mode5
+			dashboardsMode := opts.Cfg.UnifiedStorage["dashboards.dashboard.grafana.app"].DualWriterMode
+			disableDashboardsFallback = dashboardsMode == grafanarest.Mode5
+		}
+
 		// Used to get the folder stats
 		client = federated.NewFederatedClient(
 			client, // The original
 			legacysql.NewDatabaseProvider(opts.DB),
+			disableDashboardsFallback,
+			disableFoldersFallback,
 		)
 	}
 

pkg/storage/unified/federated/client.go

@@ -10,11 +10,13 @@ import (
 	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
 )
 
-func NewFederatedClient(base resource.ResourceClient, sql legacysql.LegacyDatabaseProvider) resource.ResourceClient {
+func NewFederatedClient(base resource.ResourceClient, sql legacysql.LegacyDatabaseProvider, disableDashboardsFallback bool, disableFoldersFallback bool) resource.ResourceClient {
 	return &federatedClient{
 		ResourceClient: base,
 		stats: &LegacyStatsGetter{
-			SQL: sql,
+			SQL:                          sql,
+			DisableSQLFallbackDashboards: disableDashboardsFallback,
+			DisableSQLFallbackFolders:    disableFoldersFallback,
 		},
 	}
 }

pkg/storage/unified/federated/stats.go

@@ -13,7 +13,9 @@ import (
 
 // Read stats from legacy SQL
 type LegacyStatsGetter struct {
-	SQL legacysql.LegacyDatabaseProvider
+	SQL                          legacysql.LegacyDatabaseProvider
+	DisableSQLFallbackDashboards bool
+	DisableSQLFallbackFolders    bool
 }
 
 func (s *LegacyStatsGetter) GetStats(ctx context.Context, in *resourcepb.ResourceStatsRequest) (*resourcepb.ResourceStatsResponse, error) {
@@ -64,15 +66,19 @@ func (s *LegacyStatsGetter) GetStats(ctx context.Context, in *resourcepb.Resourc
 		}
 
 		// Legacy dashboard table
-		err = fn("dashboard", "org_id=? AND folder_uid=? AND is_folder=false", group, "dashboards", true)
+		if !s.DisableSQLFallbackDashboards {
-		if err != nil {
+			err = fn("dashboard", "org_id=? AND folder_uid=? AND is_folder=false", group, "dashboards", true)
-			return err
+			if err != nil {
+				return err
+			}
 		}
 
 		// Legacy folder table
-		err = fn("folder", "org_id=? AND parent_uid=?", group, "folders", true)
+		if !s.DisableSQLFallbackFolders {
-		if err != nil {
+			err = fn("folder", "org_id=? AND parent_uid=?", group, "folders", true)
-			return err
+			if err != nil {
+				return err
+			}
 		}
 
 		// Legacy library_elements table

pkg/storage/unified/federated/stats_test.go

@@ -1,4 +1,4 @@
-package federatedtests
+package federated
 
 import (
 	"context"
@@ -22,7 +22,6 @@ import (
 	"github.com/grafana/grafana/pkg/services/tag/tagimpl"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/storage/legacysql"
-	"github.com/grafana/grafana/pkg/storage/unified/federated"
 	"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
 	"github.com/grafana/grafana/pkg/tests/testsuite"
 	"github.com/grafana/grafana/pkg/util/testutil"
@@ -32,7 +31,6 @@ func TestMain(m *testing.M) {
 	testsuite.Run(m)
 }
 
-// tests stats are correctly reported from legacy tables
 func TestIntegrationDirectSQLStats(t *testing.T) {
 	testutil.SkipIntegrationTestInShortMode(t)
 
@@ -44,7 +42,6 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 	fStore := folderimpl.ProvideStore(db)
 	tempUser := &user.SignedInUser{UserID: 1, OrgID: 1, Permissions: map[int64]map[string][]string{}}
 
-	// legacy expects the folder to be in both the dashboards and folder tables
 	folder1UID := "test1"
 	now := time.Now()
 	dashFolder1 := dashboards.NewDashboardFolder("test1")
@@ -80,7 +77,6 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 	_, err = fStore.Create(ctx, folder.CreateFolderCommand{Title: "test2", UID: folder2UID, OrgID: 1, ParentUID: folder1UID, SignedInUser: tempUser})
 	require.NoError(t, err)
 
-	// create an alert rule inside of folder test2
 	ruleStore := ngalertstore.SetupStoreForTesting(t, db)
 	_, err = ruleStore.InsertAlertRules(context.Background(), ngmodels.NewUserUID(tempUser), []ngmodels.AlertRule{
 		{
@@ -108,7 +104,6 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 		}})
 	require.NoError(t, err)
 
-	// finally, create dashboard inside of test1
 	_, err = dashStore.SaveDashboard(ctx, dashboards.SaveDashboardCommand{
 		Dashboard: simplejson.New(),
 		FolderUID: folder1UID,
@@ -116,7 +111,7 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	store := &federated.LegacyStatsGetter{
+	store := &LegacyStatsGetter{
 		SQL: legacysql.NewDatabaseProvider(db),
 	}
 
@@ -153,6 +148,105 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 		]`, string(jj))
 	})
 
+	// New tests to verify per-resource fallback disabling
+	t.Run("GetStatsForFolder1_DisableDashboardsFallback", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = request.WithNamespace(ctx, "default")
+
+		store := &LegacyStatsGetter{
+			SQL:                          legacysql.NewDatabaseProvider(db),
+			DisableSQLFallbackDashboards: true,
+			DisableSQLFallbackFolders:    false,
+		}
+
+		stats, err := store.GetStats(ctx, &resourcepb.ResourceStatsRequest{
+			Namespace: "default",
+			Folder:    folder1UID,
+		})
+		require.NoError(t, err)
+
+		var hasDashboards, hasFolders bool
+		for _, s := range stats.Stats {
+			if s.Resource == "dashboards" {
+				hasDashboards = true
+			}
+			if s.Resource == "folders" {
+				hasFolders = true
+				require.EqualValues(t, 1, s.Count)
+			}
+		}
+		require.False(t, hasDashboards, "dashboards stats should be disabled")
+		require.True(t, hasFolders, "folders stats should be present")
+	})
+
+	t.Run("GetStatsForFolder1_DisableFoldersFallback", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = request.WithNamespace(ctx, "default")
+
+		store := &LegacyStatsGetter{
+			SQL:                          legacysql.NewDatabaseProvider(db),
+			DisableSQLFallbackDashboards: false,
+			DisableSQLFallbackFolders:    true,
+		}
+
+		stats, err := store.GetStats(ctx, &resourcepb.ResourceStatsRequest{
+			Namespace: "default",
+			Folder:    folder1UID,
+		})
+		require.NoError(t, err)
+
+		var hasDashboards, hasFolders bool
+		for _, s := range stats.Stats {
+			if s.Resource == "dashboards" {
+				hasDashboards = true
+				require.EqualValues(t, 1, s.Count)
+			}
+			if s.Resource == "folders" {
+				hasFolders = true
+			}
+		}
+		require.True(t, hasDashboards, "dashboards stats should be present")
+		require.False(t, hasFolders, "folders stats should be disabled")
+	})
+
+	t.Run("GetStatsForFolder1_DisableBothFallbacks", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = request.WithNamespace(ctx, "default")
+
+		store := &LegacyStatsGetter{
+			SQL:                          legacysql.NewDatabaseProvider(db),
+			DisableSQLFallbackDashboards: true,
+			DisableSQLFallbackFolders:    true,
+		}
+
+		stats, err := store.GetStats(ctx, &resourcepb.ResourceStatsRequest{
+			Namespace: "default",
+			Folder:    folder1UID,
+		})
+		require.NoError(t, err)
+
+		var hasDashboards, hasFolders bool
+		var hasAlertRules, hasLibrary bool
+		for _, s := range stats.Stats {
+			if s.Resource == "dashboards" {
+				hasDashboards = true
+			}
+			if s.Resource == "folders" {
+				hasFolders = true
+			}
+			if s.Resource == "alertrules" {
+				hasAlertRules = true
+			}
+			if s.Resource == "library_elements" {
+				hasLibrary = true
+			}
+		}
+		require.False(t, hasDashboards, "dashboards stats should be disabled")
+		require.False(t, hasFolders, "folders stats should be disabled")
+		require.True(t, hasAlertRules, "alertrules should still be present")
+		require.True(t, hasLibrary, "library_elements should still be present")
+	})
+
 	t.Run("GetStatsForFolder2", func(t *testing.T) {
 		ctx := context.Background()
 		ctx = request.WithNamespace(ctx, "default")
@@ -178,7 +272,7 @@ func TestIntegrationDirectSQLStats(t *testing.T) {
 				"group": "sql-fallback",
 				"resource": "folders"
 			},
-						{
+			{
 				"group": "sql-fallback",
 				"resource": "library_elements"
 			}