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

5 months ago · 8bb24bc7b3
parent f51b58488c
commit 8bb24bc7b3
4 changed files with 304 additions and 2 deletions
--- a/pkg/storage/unified/client.go
+++ b/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 == "" {
--- a/pkg/storage/unified/federated/client.go
+++ b/pkg/storage/unified/federated/client.go
@ -0,0 +1,45 @@
+package federated
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+
+	"github.com/grafana/grafana/pkg/storage/legacysql"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
+)
+
+func NewFederatedClient(base resource.ResourceClient, sql legacysql.LegacyDatabaseProvider) resource.ResourceClient {
+	return &federatedClient{
+		ResourceClient: base,
+		stats: &LegacyStatsGetter{
+			SQL: sql,
+		},
+	}
+}
+
+type federatedClient struct {
+	resource.ResourceClient
+
+	// Local DB for folder stats query
+	stats *LegacyStatsGetter
+}
+
+// Get the resource stats
+func (s *federatedClient) GetStats(ctx context.Context, in *resource.ResourceStatsRequest, opts ...grpc.CallOption) (*resource.ResourceStatsResponse, error) {
+	rsp, err := s.ResourceClient.GetStats(ctx, in, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	// When folder stats are requested -- join in the legacy values
+	if in.Folder != "" {
+		more, err := s.stats.GetStats(ctx, in)
+		if err != nil {
+			return rsp, err
+		}
+		rsp.Stats = append(rsp.Stats, more.Stats...)
+	}
+
+	return rsp, err
+}
--- a/pkg/storage/unified/federated/stats.go
+++ b/pkg/storage/unified/federated/stats.go
@ -0,0 +1,70 @@
+package federated
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/grafana/authlib/claims"
+	"github.com/grafana/grafana/pkg/services/sqlstore"
+	"github.com/grafana/grafana/pkg/storage/legacysql"
+	"github.com/grafana/grafana/pkg/storage/unified/resource"
+)
+
+// Read stats from legacy SQL
+type LegacyStatsGetter struct {
+	SQL legacysql.LegacyDatabaseProvider
+}
+
+func (s *LegacyStatsGetter) GetStats(ctx context.Context, in *resource.ResourceStatsRequest) (*resource.ResourceStatsResponse, error) {
+	info, err := claims.ParseNamespace(in.Namespace)
+	if err != nil {
+		return nil, fmt.Errorf("unable to read namespace")
+	}
+	if info.OrgID == 0 {
+		return nil, fmt.Errorf("invalid OrgID found in namespace")
+	}
+
+	helper, err := s.SQL(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	rsp := &resource.ResourceStatsResponse{}
+	err = helper.DB.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
+		fn := func(table, where, g, r string) error {
+			count, err := sess.Table(helper.Table(table)).Where(where, info.OrgID, in.Folder).Count()
+			if err != nil {
+				return err
+			}
+			rsp.Stats = append(rsp.Stats, &resource.ResourceStatsResponse_Stats{
+				Group:    g, // all legacy for now
+				Resource: r,
+				Count:    count,
+			})
+			return nil
+		}
+		// Indicate that this came from the SQL tables
+		group := "sql-fallback"
+
+		// Legacy alert rule table
+		err = fn("alert_rule", "org_id=? AND dashboard_uid=?", group, "alertrules")
+		if err != nil {
+			return err
+		}
+
+		// Legacy dashboard table
+		err = fn("dashboard", "org_id=? AND folder_uid=?", group, "dashboards")
+		if err != nil {
+			return err
+		}
+
+		// Legacy folder table
+		err = fn("folder", "org_id=? AND parent_uid=?", group, "folders")
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+
+	return rsp, err
+}
--- a/pkg/storage/unified/federated/stats_test.go
+++ b/pkg/storage/unified/federated/stats_test.go
@ -0,0 +1,167 @@
+package federated
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"k8s.io/apiserver/pkg/endpoints/request"
+
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/expr"
+	"github.com/grafana/grafana/pkg/infra/db"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/accesscontrol/actest"
+	"github.com/grafana/grafana/pkg/services/accesscontrol/ossaccesscontrol/testutil"
+	"github.com/grafana/grafana/pkg/services/dashboards"
+	"github.com/grafana/grafana/pkg/services/dashboards/database"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/services/folder"
+	"github.com/grafana/grafana/pkg/services/folder/folderimpl"
+	"github.com/grafana/grafana/pkg/services/guardian"
+	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
+	ngalertstore "github.com/grafana/grafana/pkg/services/ngalert/store"
+	"github.com/grafana/grafana/pkg/services/quota/quotatest"
+	"github.com/grafana/grafana/pkg/services/supportbundles/supportbundlestest"
+	"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/resource"
+	"github.com/grafana/grafana/pkg/tests/testsuite"
+)
+
+func TestMain(m *testing.M) {
+	testsuite.Run(m)
+}
+
+func TestDirectSQLStats(t *testing.T) {
+	db, cfg := db.InitTestDBWithCfg(t)
+	ctx := context.Background()
+
+	dashStore, err := database.ProvideDashboardStore(db, cfg, featuremgmt.WithFeatures(), tagimpl.ProvideService(db), quotatest.New(false, nil))
+	require.NoError(t, err)
+	fakeGuardian := &guardian.FakeDashboardGuardian{
+		CanSaveValue: true,
+		CanEditUIDs:  []string{},
+		CanViewUIDs:  []string{},
+	}
+	guardian.MockDashboardGuardian(fakeGuardian)
+	folderPermissions, err := testutil.ProvideFolderPermissions(featuremgmt.WithFeatures(), cfg, db)
+	require.NoError(t, err)
+	fStore := folderimpl.ProvideStore(db)
+	folderSvc := folderimpl.ProvideService(fStore, actest.FakeAccessControl{ExpectedEvaluate: true}, bus.ProvideBus(tracing.InitializeTracerForTest()), dashStore,
+		folderimpl.ProvideDashboardFolderStore(db), db, featuremgmt.WithFeatures(), cfg, folderPermissions, supportbundlestest.NewFakeBundleService(), nil, tracing.InitializeTracerForTest())
+
+	// create parent folder
+
+	tempUser := &user.SignedInUser{UserID: 1, OrgID: 1, Permissions: map[int64]map[string][]string{}}
+
+	// create folders - test2 is nested in test1
+	folder1UID := "test1"
+	now := time.Now()
+	_, err = folderSvc.Create(ctx, &folder.CreateFolderCommand{Title: "test1", UID: folder1UID, OrgID: 1, SignedInUser: tempUser})
+	require.NoError(t, err)
+	folder2UID := "test2"
+	_, err = folderSvc.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.AlertRule{
+		{
+			DashboardUID: &folder2UID,
+			UID:          "test",
+			Title:        "test",
+			OrgID:        1,
+			Data: []ngmodels.AlertQuery{
+				{
+					RefID:         "A",
+					Model:         json.RawMessage("{}"),
+					DatasourceUID: expr.DatasourceUID,
+					RelativeTimeRange: ngmodels.RelativeTimeRange{
+						From: ngmodels.Duration(60),
+						To:   ngmodels.Duration(0),
+					},
+				},
+			},
+			Condition:       "ok",
+			Updated:         now,
+			NamespaceUID:    "test",
+			ExecErrState:    ngmodels.ExecutionErrorState(ngmodels.Alerting),
+			NoDataState:     ngmodels.Alerting,
+			IntervalSeconds: 60,
+		}})
+	require.NoError(t, err)
+
+	// finally, create dashboard inside of test1
+	_, err = dashStore.SaveDashboard(ctx, dashboards.SaveDashboardCommand{
+		Dashboard: simplejson.New(),
+		FolderUID: folder1UID,
+		OrgID:     1,
+	})
+	require.NoError(t, err)
+
+	store := &LegacyStatsGetter{
+		SQL: legacysql.NewDatabaseProvider(db),
+	}
+
+	t.Run("GetStatsForFolder1", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = request.WithNamespace(ctx, "default")
+
+		stats, err := store.GetStats(ctx, &resource.ResourceStatsRequest{
+			Namespace: "default",
+			Folder:    folder1UID,
+		})
+		require.NoError(t, err)
+
+		jj, _ := json.MarshalIndent(stats.Stats, "", "  ")
+		require.JSONEq(t, `[
+			{
+				"group": "sql-fallback",
+				"resource": "alertrules"
+			},
+			{
+				"group": "sql-fallback",
+				"resource": "dashboards",
+				"count": 1
+			},
+			{
+				"group": "sql-fallback",
+				"resource": "folders",
+				"count": 1
+			}
+		]`, string(jj))
+	})
+
+	t.Run("GetStatsForFolder2", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = request.WithNamespace(ctx, "default")
+
+		stats, err := store.GetStats(ctx, &resource.ResourceStatsRequest{
+			Namespace: "default",
+			Folder:    folder2UID,
+		})
+		require.NoError(t, err)
+
+		jj, _ := json.MarshalIndent(stats.Stats, "", "  ")
+		require.JSONEq(t, `[
+			{
+				"group": "sql-fallback",
+				"resource": "alertrules",
+				"count": 1
+			},
+			{
+				"group": "sql-fallback",
+				"resource": "dashboards"
+			},
+			{
+				"group": "sql-fallback",
+				"resource": "folders"
+			}
+		]`, string(jj))
+	})
+}