From 8bb24bc7b3d5d9e471b4609538d89e9063a0cdef Mon Sep 17 00:00:00 2001 From: Ryan McKinley Date: Wed, 11 Dec 2024 14:46:31 +0300 Subject: [PATCH] Unified: Add client-side stats federation to support folders (#97778) --- pkg/storage/unified/client.go | 24 ++- pkg/storage/unified/federated/client.go | 45 ++++++ pkg/storage/unified/federated/stats.go | 70 ++++++++ pkg/storage/unified/federated/stats_test.go | 167 ++++++++++++++++++++ 4 files changed, 304 insertions(+), 2 deletions(-) create mode 100644 pkg/storage/unified/federated/client.go create mode 100644 pkg/storage/unified/federated/stats.go create mode 100644 pkg/storage/unified/federated/stats_test.go diff --git a/pkg/storage/unified/client.go b/pkg/storage/unified/client.go index e3c8a13ddbd..3c45d2135e4 100644 --- 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 == "" { diff --git a/pkg/storage/unified/federated/client.go b/pkg/storage/unified/federated/client.go new file mode 100644 index 00000000000..666ef7d3625 --- /dev/null +++ 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 +} diff --git a/pkg/storage/unified/federated/stats.go b/pkg/storage/unified/federated/stats.go new file mode 100644 index 00000000000..8e5f7c0033c --- /dev/null +++ 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 +} diff --git a/pkg/storage/unified/federated/stats_test.go b/pkg/storage/unified/federated/stats_test.go new file mode 100644 index 00000000000..56181303632 --- /dev/null +++ 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)) + }) +}