Unified Storage Search: Add sprinkles (#97820)

* Wire up sprinkles to oss and enterprise. Fetching sprinkles not implemented yet.

* Adds wireset for initializing document builders. Had to init it when creating the service to avoid cyclical imports.

* updates to int64 for stats

* adds config for sprinklesApiServer and gets sprinkles from there when its present

* add comment for later

* adds feature toggle for sprinkles. returns empty results when flag not enabled.

* adds unified storage config setting for sprinkles apiserver page limit

* fixes bug where dashboard uid was not getting set

* when creating dashboard summary, use metadata.name as the dashboard uid

* cleans up wire. use existing oss and enterprise sets to generate doc builders

* remove old wireset

* fix linter - adds missing arg for doc builders

* update dashboard stats in tests

* updates test-data dashboards

* log a warning instead of returning an error if we can't get sprinkles for a namespace

* dont read uid from dashboard json

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -220,6 +220,7 @@ export interface FeatureToggles {
   useSessionStorageForRedirection?: boolean;
   rolePickerDrawer?: boolean;
   unifiedStorageSearch?: boolean;
+  unifiedStorageSearchSprinkles?: boolean;
   pluginsSriChecks?: boolean;
   unifiedStorageBigObjectsSupport?: boolean;
   timeRangeProvider?: boolean;

pkg/server/module_server.go

@@ -131,7 +131,11 @@ func (s *ModuleServer) Run() error {
 	//}
 
 	m.RegisterModule(modules.StorageServer, func() (services.Service, error) {
-		return sql.ProvideUnifiedStorageGrpcService(s.cfg, s.features, nil, s.log, nil)
+		docBuilders, err := InitializeDocumentBuilders(s.cfg)
+		if err != nil {
+			return nil, err
+		}
+		return sql.ProvideUnifiedStorageGrpcService(s.cfg, s.features, nil, s.log, nil, docBuilders)
 	})
 
 	m.RegisterModule(modules.ZanzanaServer, func() (services.Service, error) {

pkg/server/wire.go

@@ -8,6 +8,7 @@ package server
 
 import (
 	"github.com/google/wire"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
 
 	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
 
@@ -229,7 +230,6 @@ var wireBasicSet = wire.NewSet(
 	wire.Bind(new(login.AuthInfoService), new(*authinfoimpl.Service)),
 	authinfoimpl.ProvideStore,
 	datasourceproxy.ProvideService,
-	unifiedsearch.ProvideDocumentBuilders,
 	search.ProvideService,
 	searchV2.ProvideService,
 	searchV2.ProvideSearchHTTPService,
@@ -475,3 +475,8 @@ func InitializeAPIServerFactory() (standalone.APIServerFactory, error) {
 	wire.Build(wireExtsStandaloneAPIServerSet)
 	return &standalone.NoOpAPIServerFactory{}, nil // Wire will replace this with a real interface
 }
+
+func InitializeDocumentBuilders(cfg *setting.Cfg) (resource.DocumentBuilderSupplier, error) {
+	wire.Build(wireExtsSet)
+	return &unifiedsearch.StandardDocumentBuilders{}, nil
+}

pkg/server/wireexts_oss.go

@@ -6,6 +6,7 @@ package server
 
 import (
 	"github.com/google/wire"
+	search2 "github.com/grafana/grafana/pkg/storage/unified/search"
 
 	"github.com/grafana/grafana/pkg/infra/metrics"
 	"github.com/grafana/grafana/pkg/plugins"
@@ -106,6 +107,9 @@ var wireExtsBasicSet = wire.NewSet(
 	wire.Bind(new(auth.IDSigner), new(*idimpl.LocalSigner)),
 	manager.ProvideInstaller,
 	wire.Bind(new(plugins.Installer), new(*manager.PluginInstaller)),
+	search2.ProvideDashboardStats,
+	wire.Bind(new(search2.DashboardStats), new(*search2.OssDashboardStats)),
+	search2.ProvideDocumentBuilders,
 )
 
 var wireExtsSet = wire.NewSet(

pkg/services/featuremgmt/registry.go

@@ -1525,6 +1525,14 @@ var (
 			HideFromDocs:      true,
 			HideFromAdminPage: true,
 		},
+		{
+			Name:              "unifiedStorageSearchSprinkles",
+			Description:       "Enable sprinkles on unified storage search",
+			Stage:             FeatureStageExperimental,
+			Owner:             grafanaSearchAndStorageSquad,
+			HideFromDocs:      true,
+			HideFromAdminPage: true,
+		},
 		{
 			Name:        "pluginsSriChecks",
 			Description: "Enables SRI checks for plugin assets",

pkg/services/featuremgmt/toggles_gen.csv

@@ -201,6 +201,7 @@ improvedExternalSessionHandling,experimental,@grafana/identity-access-team,false
 useSessionStorageForRedirection,GA,@grafana/identity-access-team,false,false,false
 rolePickerDrawer,experimental,@grafana/identity-access-team,false,false,false
 unifiedStorageSearch,experimental,@grafana/search-and-storage,false,false,false
+unifiedStorageSearchSprinkles,experimental,@grafana/search-and-storage,false,false,false
 pluginsSriChecks,experimental,@grafana/plugins-platform-backend,false,false,false
 unifiedStorageBigObjectsSupport,experimental,@grafana/search-and-storage,false,false,false
 timeRangeProvider,experimental,@grafana/grafana-frontend-platform,false,false,false

pkg/services/featuremgmt/toggles_gen.go

@@ -815,6 +815,10 @@ const (
 	// Enable unified storage search
 	FlagUnifiedStorageSearch = "unifiedStorageSearch"
 
+	// FlagUnifiedStorageSearchSprinkles
+	// Enable sprinkles on unified storage search
+	FlagUnifiedStorageSearchSprinkles = "unifiedStorageSearchSprinkles"
+
 	// FlagPluginsSriChecks
 	// Enables SRI checks for plugin assets
 	FlagPluginsSriChecks = "pluginsSriChecks"

pkg/services/featuremgmt/toggles_gen.json

@@ -3618,6 +3618,20 @@
         "hideFromDocs": true
       }
     },
+    {
+      "metadata": {
+        "name": "unifiedStorageSearchSprinkles",
+        "resourceVersion": "1734123247356",
+        "creationTimestamp": "2024-12-13T20:54:07Z"
+      },
+      "spec": {
+        "description": "Enable sprinkles on unified storage search",
+        "stage": "experimental",
+        "codeowner": "@grafana/search-and-storage",
+        "hideFromAdminPage": true,
+        "hideFromDocs": true
+      }
+    },
     {
       "metadata": {
         "name": "useSeessionStorageForRedirection",

pkg/setting/setting.go

@@ -523,12 +523,14 @@ type Cfg struct {
 	ShortLinkExpiration int
 
 	// Unified Storage
-	UnifiedStorage     map[string]UnifiedStorageConfig
-	IndexPath          string
-	IndexWorkers       int
-	IndexMaxBatchSize  int
-	IndexFileThreshold int
-	IndexMinCount      int
+	UnifiedStorage              map[string]UnifiedStorageConfig
+	IndexPath                   string
+	IndexWorkers                int
+	IndexMaxBatchSize           int
+	IndexFileThreshold          int
+	IndexMinCount               int
+	SprinklesApiServer          string
+	SprinklesApiServerPageLimit int
 }
 
 type UnifiedStorageConfig struct {

pkg/setting/setting_unified_storage.go

@@ -52,4 +52,6 @@ func (cfg *Cfg) setUnifiedStorageConfig() {
 	cfg.IndexMaxBatchSize = section.Key("index_max_batch_size").MustInt(100)
 	cfg.IndexFileThreshold = section.Key("index_file_threshold").MustInt(10)
 	cfg.IndexMinCount = section.Key("index_min_count").MustInt(1)
+	cfg.SprinklesApiServer = section.Key("sprinkles_api_server").String()
+	cfg.SprinklesApiServerPageLimit = section.Key("sprinkles_api_server_page_limit").MustInt(100)
 }

pkg/storage/unified/search/bleve_test.go

@@ -49,7 +49,7 @@ func TestBleveBackend(t *testing.T) {
 			return &DashboardDocumentBuilder{
 				Namespace:        namespace,
 				Blob:             blob,
-				Stats:            NewDashboardStatsLookup(nil), // empty stats
+				Stats:            make(map[string]map[string]int64), // empty stats
 				DatasourceLookup: dashboard.CreateDatasourceLookup([]*dashboard.DatasourceQueryResult{{}}),
 			}, nil
 		})

pkg/storage/unified/search/dashboard.go

@@ -75,7 +75,7 @@ func DashboardBuilder(namespaced resource.NamespacedDocumentSupplier) (resource.
 			return &DashboardDocumentBuilder{
 				Namespace:        namespace,
 				Blob:             blob,
-				Stats:            NewDashboardStatsLookup(nil),
+				Stats:            nil,
 				DatasourceLookup: dashboard.CreateDatasourceLookup([]*dashboard.DatasourceQueryResult{
 					// empty values (does not resolve anything)
 				}),
@@ -94,8 +94,8 @@ type DashboardDocumentBuilder struct {
 	Namespace string
 
 	// Cached stats for this namespace
-	// TODO, load this from apiserver request
-	Stats DashboardStatsLookup
+	// maps dashboard UID to stats
+	Stats map[string]map[string]int64
 
 	// data source lookup
 	DatasourceLookup dashboard.DatasourceLookup
@@ -104,17 +104,12 @@ type DashboardDocumentBuilder struct {
 	Blob resource.BlobSupport
 }
 
-type DashboardStatsLookup = func(ctx context.Context, uid string) map[string]int64
-
-func NewDashboardStatsLookup(stats map[string]map[string]int64) DashboardStatsLookup {
-	return func(ctx context.Context, uid string) map[string]int64 {
-		if stats == nil {
-			return nil
-		}
-		return stats[uid]
-	}
+type DashboardStats interface {
+	GetStats(ctx context.Context, namespace string) (map[string]map[string]int64, error)
 }
 
+type DashboardStatsLookup = func(ctx context.Context, uid string) map[string]int64
+
 var _ resource.DocumentBuilder = &DashboardDocumentBuilder{}
 
 func (s *DashboardDocumentBuilder) BuildDocument(ctx context.Context, key *resource.ResourceKey, rv int64, value []byte) (*resource.IndexableDocument, error) {
@@ -150,6 +145,9 @@ func (s *DashboardDocumentBuilder) BuildDocument(ctx context.Context, key *resou
 		return nil, err
 	}
 
+	// metadata name is the dashboard uid
+	summary.UID = obj.GetName()
+
 	doc := resource.NewIndexableDocument(key, rv, obj)
 	doc.Title = summary.Title
 	doc.Description = summary.Description
@@ -211,8 +209,7 @@ func (s *DashboardDocumentBuilder) BuildDocument(ctx context.Context, key *resou
 	}
 
 	// Add the stats fields
-	stats := s.Stats(ctx, key.Name) // summary.UID
-	for k, v := range stats {
+	for k, v := range s.Stats[summary.UID] {
 		doc.Fields[k] = v
 	}
 

pkg/storage/unified/search/dashboard_stats.go

@@ -0,0 +1,15 @@
+package search
+
+import (
+	"context"
+)
+
+type OssDashboardStats struct{}
+
+func ProvideDashboardStats() *OssDashboardStats {
+	return &OssDashboardStats{}
+}
+
+func (s *OssDashboardStats) GetStats(ctx context.Context, namespace string) (map[string]map[string]int64, error) {
+	return nil, nil
+}

pkg/storage/unified/search/document.go

@@ -5,23 +5,25 @@ import (
 
 	"github.com/grafana/authlib/claims"
 	"github.com/grafana/grafana/pkg/infra/db"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/store/kind/dashboard"
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
 )
 
 // The default list of open source document builders
 type StandardDocumentBuilders struct {
-	sql db.DB
+	sql       db.DB
+	sprinkles DashboardStats
 }
 
 // Hooked up so wire can fill in different sprinkles
-func ProvideDocumentBuilders(sql db.DB) resource.DocumentBuilderSupplier {
-	return &StandardDocumentBuilders{sql}
+func ProvideDocumentBuilders(sql db.DB, sprinkles DashboardStats) resource.DocumentBuilderSupplier {
+	return &StandardDocumentBuilders{sql, sprinkles}
 }
 
 func (s *StandardDocumentBuilders) GetDocumentBuilders() ([]resource.DocumentBuilderInfo, error) {
 	dashboards, err := DashboardBuilder(func(ctx context.Context, namespace string, blob resource.BlobSupport) (resource.DocumentBuilder, error) {
-		stats := NewDashboardStatsLookup(nil) // empty stats
+		logger := log.New("dashboard_builder", "namespace", namespace)
 		dsinfo := []*dashboard.DatasourceQueryResult{{}}
 		ns, err := claims.ParseNamespace(namespace)
 		if err != nil && s.sql != nil {
@@ -43,6 +45,18 @@ func (s *StandardDocumentBuilders) GetDocumentBuilders() ([]resource.DocumentBui
 				dsinfo = append(dsinfo, info)
 			}
 		}
+
+		// Fetch dashboard sprinkles for the namespace
+		// This could take a while if namespace has a lot of dashboards
+		var stats map[string]map[string]int64
+		if s.sprinkles != nil {
+			stats, err = s.sprinkles.GetStats(ctx, namespace)
+			if err != nil {
+				// only log a warning. Don't need to fail the indexer if we can't get sprinkles
+				logger.Warn("Failed to get sprinkles", "error", err)
+			}
+		}
+
 		return &DashboardDocumentBuilder{
 			Namespace:        namespace,
 			Blob:             blob,

pkg/storage/unified/search/document_test.go

@@ -56,12 +56,12 @@ func TestDashboardDocumentBuilder(t *testing.T) {
 		return &DashboardDocumentBuilder{
 			Namespace: namespace,
 			Blob:      blob,
-			Stats: NewDashboardStatsLookup(map[string]map[string]int64{
+			Stats: map[string]map[string]int64{
 				"aaa": {
 					DASHBOARD_ERRORS_LAST_1_DAYS: 1,
 					DASHBOARD_ERRORS_LAST_7_DAYS: 1,
 				},
-			}),
+			},
 			DatasourceLookup: dashboard.CreateDatasourceLookup([]*dashboard.DatasourceQueryResult{{
 				Name: "TheDisplayName", // used to be the unique ID!
 				Type: "my-custom-plugin",

pkg/storage/unified/sql/service.go

@@ -20,7 +20,6 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/grafana/grafana/pkg/storage/unified/resource/grpc"
-	"github.com/grafana/grafana/pkg/storage/unified/search"
 )
 
 var (
@@ -51,6 +50,8 @@ type service struct {
 
 	log log.Logger
 	reg prometheus.Registerer
+
+	docBuilders resource.DocumentBuilderSupplier
 }
 
 func ProvideUnifiedStorageGrpcService(
@@ -59,6 +60,7 @@ func ProvideUnifiedStorageGrpcService(
 	db infraDB.DB,
 	log log.Logger,
 	reg prometheus.Registerer,
+	docBuilders resource.DocumentBuilderSupplier,
 ) (UnifiedStorageGrpcService, error) {
 	tracingCfg, err := tracing.ProvideTracingConfig(cfg)
 	if err != nil {
@@ -78,7 +80,7 @@ func ProvideUnifiedStorageGrpcService(
 
 	// FIXME: This is a temporary solution while we are migrating to the new authn interceptor
 	// grpcutils.NewGrpcAuthenticator should be used instead.
-	authn, err := grpcutils.NewGrpcAuthenticatorWithFallback(cfg, prometheus.DefaultRegisterer, tracing, &grpc.Authenticator{})
+	authn, err := grpcutils.NewGrpcAuthenticatorWithFallback(cfg, reg, tracing, &grpc.Authenticator{})
 	if err != nil {
 		return nil, err
 	}
@@ -92,6 +94,7 @@ func ProvideUnifiedStorageGrpcService(
 		db:            db,
 		log:           log,
 		reg:           reg,
+		docBuilders:   docBuilders,
 	}
 
 	// This will be used when running as a dskit service
@@ -106,11 +109,7 @@ func (s *service) start(ctx context.Context) error {
 		return err
 	}
 
-	// TODO, for standalone this will need to be started from enterprise
-	// Connecting to the correct remote services (cloudconfig for DS info and usage stats)
-	docs := search.ProvideDocumentBuilders(nil)
-
-	server, err := NewResourceServer(ctx, s.db, s.cfg, s.features, docs, s.tracing, s.reg, authzClient)
+	server, err := NewResourceServer(ctx, s.db, s.cfg, s.features, s.docBuilders, s.tracing, s.reg, authzClient)
 	if err != nil {
 		return err
 	}

pkg/storage/unified/sql/test/integration_test.go

@@ -366,7 +366,7 @@ func TestClientServer(t *testing.T) {
 
 	features := featuremgmt.WithFeatures()
 
-	svc, err := sql.ProvideUnifiedStorageGrpcService(cfg, features, dbstore, nil, prometheus.NewPedanticRegistry())
+	svc, err := sql.ProvideUnifiedStorageGrpcService(cfg, features, dbstore, nil, prometheus.NewPedanticRegistry(), nil)
 	require.NoError(t, err)
 	var client resource.ResourceStoreClient
 

pkg/tests/testinfra/testinfra.go

@@ -104,7 +104,7 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 	var storage sql.UnifiedStorageGrpcService
 	if runstore {
 		storage, err = sql.ProvideUnifiedStorageGrpcService(env.Cfg, env.FeatureToggles, env.SQLStore,
-			env.Cfg.Logger, prometheus.NewPedanticRegistry())
+			env.Cfg.Logger, prometheus.NewPedanticRegistry(), nil)
 		require.NoError(t, err)
 		ctx := context.Background()
 		err = storage.StartAsync(ctx)