Unified: Add client-side stats federation to support folders (#97778)

pkg/storage/unified/client.go

@@ -20,6 +20,8 @@ import (
 	"github.com/grafana/grafana/pkg/services/authz"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/storage/legacysql"
+	"github.com/grafana/grafana/pkg/storage/unified/federated"
 	"github.com/grafana/grafana/pkg/storage/unified/resource"
 	"github.com/grafana/grafana/pkg/storage/unified/sql"
 )
@@ -38,14 +40,32 @@ func ProvideUnifiedStorageClient(
 ) (resource.ResourceClient, error) {
 	// See: apiserver.ApplyGrafanaConfig(cfg, features, o)
 	apiserverCfg := cfg.SectionWithEnvOverrides("grafana-apiserver")
-	opts := options.StorageOptions{
+	client, err := newClient(options.StorageOptions{
 		StorageType:  options.StorageType(apiserverCfg.Key("storage_type").MustString(string(options.StorageTypeLegacy))),
 		DataPath:     apiserverCfg.Key("storage_path").MustString(filepath.Join(cfg.DataPath, "grafana-apiserver")),
 		Address:      apiserverCfg.Key("address").MustString(""), // client address
 		BlobStoreURL: apiserverCfg.Key("blob_url").MustString(""),
+	}, cfg, features, db, tracer, reg, authzc, docs)
+	if err == nil {
+		// Used to get the folder stats
+		client = federated.NewFederatedClient(
+			client, // The original
+			legacysql.NewDatabaseProvider(db),
+		)
 	}
-	ctx := context.Background()
+	return client, err
+}
 
+func newClient(opts options.StorageOptions,
+	cfg *setting.Cfg,
+	features featuremgmt.FeatureToggles,
+	db infraDB.DB,
+	tracer tracing.Tracer,
+	reg prometheus.Registerer,
+	authzc authz.Client,
+	docs resource.DocumentBuilderSupplier,
+) (resource.ResourceClient, error) {
+	ctx := context.Background()
 	switch opts.StorageType {
 	case options.StorageTypeFile:
 		if opts.DataPath == "" {