Alerting: Improve ASH Loki query efficiency by including folderUID (#113322)

* Alerting: Improve ASH Loki query efficiency by including folderUID

Previously, the folderUID label was only included when ruleUID was not specified
 and the user did not have full alert rule read permissions.

To improve ASH Loki query efficiency, this PR includes the folderUID in the ASH
Loki query when ruleUID is specified, even if the user has full alert rule read
permissions.

Some non-obvious considerations:
- The naive implementation of just including the current folder UID would have
the unintended side-effect of no longer returning history after a rule is moved
 between folders.
- The previous implementation made the trade-off of only checking RBAC on the
current folder, including any history from old folders that may exist.

To solve both of the above, we make an extra query to the database to check the
alert rule's previous versions so we can include any old folderUIDs, checking
RBAC at the same time.

The querying and inclusion of history from old folders is done best-effort, any
issues that might arise are logged and ignored so as not to prevent the current
folder history.

* Fix merge conflicts

* Reduce scanning on GetAlertRuleVersionFolders by grouping in query

pkg/services/ngalert/state/historian/annotation.go

@@ -44,6 +44,7 @@ type AnnotationBackend struct {
 type RuleStore interface {
 	GetAlertRuleByUID(ctx context.Context, query *ngmodels.GetAlertRuleByUIDQuery) (*ngmodels.AlertRule, error)
 	GetUserVisibleNamespaces(ctx context.Context, orgID int64, user identity.Requester) (map[string]*folder.Folder, error)
+	GetAlertRuleVersionFolders(ctx context.Context, orgID int64, guid string) ([]string, error)
 }
 
 type AnnotationStore interface {

pkg/services/ngalert/state/historian/loki.go

@@ -505,30 +505,25 @@ func (h *RemoteLokiBackend) getFolderUIDsForFilter(ctx context.Context, query mo
 	if err != nil {
 		return nil, err
 	}
-	if bypass { // if user has access to all rules and folder, remove filter
+
+	if query.RuleUID != "" {
+		return h.getFolderUIDsForRuleFilter(ctx, query, bypass)
+	}
+
+	// If the query has no rule filter, we need to return all folder UIDs the user has access to.
+	// For a user with access to all rules and folders, the full list of folders will likely be too large to be an
+	// effective optimization in Loki, so we skip folderUID filtering entirely in that case.
+	if bypass {
 		return nil, nil
 	}
-	// if there is a filter by rule UID, find that rule UID and make sure that user has access to it.
-	if query.RuleUID != "" {
-		rule, err := h.ruleStore.GetAlertRuleByUID(ctx, &models.GetAlertRuleByUIDQuery{
-			UID:   query.RuleUID,
-			OrgID: query.OrgID,
-		})
-		if err != nil {
-			return nil, fmt.Errorf("failed to fetch alert rule by UID: %w", err)
-		}
-		if rule == nil {
-			return nil, models.ErrAlertRuleNotFound
-		}
-		return nil, h.ac.AuthorizeAccessInFolder(ctx, query.SignedInUser, rule)
-	}
-	// if no filter, then we need to get all namespaces user has access to
+
+	// All folders the user has access to.
 	folders, err := h.ruleStore.GetUserVisibleNamespaces(ctx, query.OrgID, query.SignedInUser)
 	if err != nil {
 		return nil, fmt.Errorf("failed to fetch folders that user can access: %w", err)
 	}
 	uids := make([]string, 0, len(folders))
-	// now keep only UIDs of folder in which user can read rules.
+	// Keep only UIDs of folder in which user can read rules.
 	for _, f := range folders {
 		hasAccess, err := h.ac.HasAccessInFolder(ctx, query.SignedInUser, models.NewNamespace(f))
 		if err != nil {
@@ -545,3 +540,71 @@ func (h *RemoteLokiBackend) getFolderUIDsForFilter(ctx context.Context, query mo
 	sort.Strings(uids)
 	return uids, nil
 }
+
+func (h *RemoteLokiBackend) getFolderUIDsForRuleFilter(ctx context.Context, query models.HistoryQuery, canReadAll bool) ([]string, error) {
+	rule, err := h.ruleStore.GetAlertRuleByUID(ctx, &models.GetAlertRuleByUIDQuery{
+		UID:   query.RuleUID,
+		OrgID: query.OrgID,
+	})
+	if err != nil {
+		if canReadAll {
+			// When the user can read all rules, filtering by folder UID is purely an optimization, so we can ignore errors here.
+			h.log.FromContext(ctx).Debug("failed to fetch alert rule by UID", "err", err)
+			return nil, nil
+		}
+		return nil, fmt.Errorf("failed to fetch alert rule by UID: %w", err)
+	}
+
+	// First, we check if the user has access to the current version of the rule. If not, we can return early.
+	// Whether we should check historical folders they might still have access to is not 100% clear, but it seems more
+	// intuitive to deny access in this case.
+	if !canReadAll {
+		if err := h.ac.AuthorizeAccessInFolder(ctx, query.SignedInUser, rule); err != nil {
+			return nil, err
+		}
+	}
+
+	// We want to return folder UIDs when possible, as it's indexed in Loki and will help with query performance.
+	// However, by just returning the current folder UID the user can lose history when a rule is moved between folders.
+	// So, we attempt to get historical folder UIDs from the rule's history.
+	historicalFolders, err := h.ruleStore.GetAlertRuleVersionFolders(ctx, rule.OrgID, rule.GUID)
+	if err != nil {
+		// Including historical folders is an edge case enhancement, better to just log the error and continue
+		// with the current folder UID.
+		h.log.FromContext(ctx).Debug("failed to include historical folder UIDs for rule", "err", err)
+	}
+
+	accessibleFolders := make([]string, 0, len(historicalFolders)+1)
+	dedup := make(map[string]struct{})
+
+	accessibleFolders = append(accessibleFolders, rule.GetNamespaceUID())
+	dedup[rule.GetNamespaceUID()] = struct{}{}
+
+	for _, folderUID := range historicalFolders {
+		if _, exists := dedup[folderUID]; exists {
+			continue
+		}
+
+		if canReadAll {
+			// If the user can read all rules, no need to check access to each folder.
+			accessibleFolders = append(accessibleFolders, folderUID)
+			continue
+		}
+
+		hasAccess, err := h.ac.HasAccessInFolder(ctx, query.SignedInUser, models.Namespace{
+			UID: folderUID,
+		})
+		if err != nil {
+			// Including historical folders is an edge case enhancement, better to just log the error and continue
+			// with the current folder UID.
+			h.log.FromContext(ctx).Debug("failed to check access to folder", "err", err, "folderUID", folderUID)
+			continue
+		}
+		if !hasAccess {
+			continue
+		}
+		accessibleFolders = append(accessibleFolders, folderUID)
+	}
+
+	return accessibleFolders, nil
+}

pkg/services/ngalert/state/historian/loki_test.go

@@ -22,6 +22,7 @@ import (
 
 	alertingInstrument "github.com/grafana/alerting/http/instrument"
 	"github.com/grafana/alerting/http/instrument/instrumenttest"
+
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -868,7 +869,8 @@ func TestGetFolderUIDsForFilter(t *testing.T) {
 			}
 			result, err := createLoki(ac).getFolderUIDsForFilter(context.Background(), models.HistoryQuery{OrgID: orgID, RuleUID: rule.UID, SignedInUser: usr})
 			assert.NoError(t, err)
-			assert.Empty(t, result)
+			assert.Len(t, result, 1)
+			assert.Contains(t, result, rule.GetNamespaceUID())
 
 			assert.Len(t, ac.Calls, 1)
 			assert.Equal(t, "CanReadAllRules", ac.Calls[0].MethodName)
@@ -893,7 +895,8 @@ func TestGetFolderUIDsForFilter(t *testing.T) {
 
 			result, err := loki.getFolderUIDsForFilter(context.Background(), models.HistoryQuery{OrgID: orgID, RuleUID: rule.UID, SignedInUser: usr})
 			assert.NoError(t, err)
-			assert.Empty(t, result)
+			assert.Len(t, result, 1)
+			assert.Contains(t, result, rule.GetNamespaceUID())
 
 			assert.Len(t, ac.Calls, 2)
 			assert.Equal(t, "CanReadAllRules", ac.Calls[0].MethodName)
@@ -916,6 +919,21 @@ func TestGetFolderUIDsForFilter(t *testing.T) {
 				require.ErrorIs(t, err, models.ErrAlertRuleNotFound)
 			})
 		})
+
+		t.Run("should return folderUID", func(t *testing.T) {
+			for _, authBypass := range []bool{true, false} {
+				t.Run(fmt.Sprintf("authBypass=%v", authBypass), func(t *testing.T) {
+					ac := &acfakes.FakeRuleService{}
+					ac.CanReadAllRulesFunc = func(ctx context.Context, requester identity.Requester) (bool, error) {
+						return authBypass, nil
+					}
+					result, err := createLoki(ac).getFolderUIDsForFilter(context.Background(), models.HistoryQuery{OrgID: orgID, RuleUID: rule.UID, SignedInUser: usr})
+					assert.NoError(t, err)
+					assert.Len(t, result, 1)
+					assert.Contains(t, result, rule.GetNamespaceUID())
+				})
+			}
+		})
 	})
 
 	t.Run("when rule UID is empty", func(t *testing.T) {
@@ -982,6 +1000,161 @@ func TestGetFolderUIDsForFilter(t *testing.T) {
 	})
 }
 
+func TestGetFolderUIDsForFilterWithHistoricalFolders(t *testing.T) {
+	// Simple history generator to avoid repetitive code in test cases.
+	simpleHistory := func(guid string) map[string][]*models.AlertRuleVersion {
+		return map[string][]*models.AlertRuleVersion{
+			guid: {
+				&models.AlertRuleVersion{AlertRule: models.RuleGen.With(models.RuleMuts.WithGUID(guid), models.RuleMuts.WithNamespaceUID("folder-current")).Generate()},
+				&models.AlertRuleVersion{AlertRule: models.RuleGen.With(models.RuleMuts.WithGUID(guid), models.RuleMuts.WithNamespaceUID("folder-historical-1")).Generate()},
+				&models.AlertRuleVersion{AlertRule: models.RuleGen.With(models.RuleMuts.WithGUID(guid), models.RuleMuts.WithNamespaceUID("folder-historical-2")).Generate()},
+				&models.AlertRuleVersion{AlertRule: models.RuleGen.With(models.RuleMuts.WithGUID(guid), models.RuleMuts.WithNamespaceUID("folder-historical-3")).Generate()},
+			},
+		}
+	}
+
+	// Helper to create simple folder access override functions.
+	canReadRulesInFolders := func(folderUids ...string) func(folderUID string) (bool, error) {
+		return func(folderUID string) (bool, error) {
+			for _, f := range folderUids {
+				if folderUID == f {
+					return true, nil
+				}
+			}
+			return false, nil
+		}
+	}
+
+	// Helper to fail historical folders query.
+	failHistoryQueryHook := func(cmd any) error {
+		q, ok := cmd.(fakes.GenericRecordedQuery)
+		if !ok {
+			return nil
+		}
+		if q.Name == "GetAlertRuleVersionFolders" {
+			return errors.New("generic error")
+		}
+		return nil
+	}
+
+	cases := []struct {
+		name string
+		// Setup.
+		existingHistory map[string][]*models.AlertRuleVersion
+		canReadAllRules bool
+		rule            *models.AlertRule
+
+		// Error overrides.
+		ruleStoreHook        func(cmd any) error
+		folderAccessOverride func(folderUID string) (bool, error)
+
+		// Expected.
+		expectedFolders []string
+	}{
+		{
+			name:            "should include historical folders when user can read all rules",
+			existingHistory: simpleHistory("guid-1"),
+			canReadAllRules: true,
+			rule:            models.RuleGen.With(models.RuleMuts.WithGUID("guid-1"), models.RuleMuts.WithNamespaceUID("folder-current")).GenerateRef(),
+			expectedFolders: []string{"folder-current", "folder-historical-1", "folder-historical-2", "folder-historical-3"},
+		},
+		{
+			name:                 "should include only authorized historical folders",
+			existingHistory:      simpleHistory("guid-1"),
+			folderAccessOverride: canReadRulesInFolders("folder-current", "folder-historical-2"),
+			rule:                 models.RuleGen.With(models.RuleMuts.WithGUID("guid-1"), models.RuleMuts.WithNamespaceUID("folder-current")).GenerateRef(),
+			expectedFolders:      []string{"folder-current", "folder-historical-2"},
+		},
+		{
+			name:            "if historical folders query fails, should return current folder",
+			existingHistory: simpleHistory("guid-1"),
+			canReadAllRules: true,
+			rule:            models.RuleGen.With(models.RuleMuts.WithGUID("guid-1"), models.RuleMuts.WithNamespaceUID("folder-current")).GenerateRef(),
+			ruleStoreHook:   failHistoryQueryHook,
+			expectedFolders: []string{"folder-current"},
+		},
+		{
+			name:                 "if historical folders query fails, should return current folder",
+			existingHistory:      simpleHistory("guid-1"),
+			folderAccessOverride: canReadRulesInFolders("folder-current", "folder-historical-2"),
+			rule:                 models.RuleGen.With(models.RuleMuts.WithGUID("guid-1"), models.RuleMuts.WithNamespaceUID("folder-current")).GenerateRef(),
+			ruleStoreHook:        failHistoryQueryHook,
+			expectedFolders:      []string{"folder-current"},
+		},
+		{
+			name:            "if access check for historical folders fails, should ignore",
+			existingHistory: simpleHistory("guid-1"),
+			rule:            models.RuleGen.With(models.RuleMuts.WithGUID("guid-1"), models.RuleMuts.WithNamespaceUID("folder-current")).GenerateRef(),
+			folderAccessOverride: func(folderUID string) (bool, error) {
+				switch folderUID {
+				case "folder-current", "folder-historical-3":
+					return true, nil
+				case "folder-historical-2":
+					return false, nil
+				case "folder-historical-1":
+					return false, errors.New("generic error")
+				}
+				return false, nil
+			},
+			expectedFolders: []string{"folder-current", "folder-historical-3"},
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Setup.
+			orgID := int64(1)
+			usr := accesscontrol.BackgroundUser("test", 1, org.RoleNone, nil)
+
+			ac := &acfakes.FakeRuleService{}
+			ac.CanReadAllRulesFunc = func(ctx context.Context, requester identity.Requester) (bool, error) {
+				return tc.canReadAllRules, nil
+			}
+			ac.AuthorizeAccessInFolderFunc = func(ctx context.Context, requester identity.Requester, namespaced models.Namespaced) error {
+				if tc.canReadAllRules {
+					return nil
+				}
+				hasAccess, err := tc.folderAccessOverride(namespaced.GetNamespaceUID())
+				if err != nil {
+					return err
+				}
+				if !hasAccess {
+					return rulesAuthz.ErrAuthorizationBase.Errorf("%w", err)
+				}
+				return nil
+			}
+			ac.HasAccessInFolderFunc = func(ctx context.Context, requester identity.Requester, namespaced models.Namespaced) (bool, error) {
+				if tc.canReadAllRules {
+					return true, nil
+				}
+				return tc.folderAccessOverride(namespaced.GetNamespaceUID())
+			}
+
+			rulesStore := fakes.NewRuleStore(t)
+			rulesStore.Rules = map[int64][]*models.AlertRule{
+				orgID: {
+					tc.rule,
+					// Add some irrelevant rules to ensure they are ignored.
+					models.RuleGen.With(models.RuleMuts.WithNamespaceUID(tc.rule.GetNamespaceUID())).GenerateRef(),
+					models.RuleGen.With(models.RuleMuts.WithNamespaceUID("irrelevant-folder")).GenerateRef(),
+				},
+			}
+			rulesStore.History = tc.existingHistory
+			if tc.ruleStoreHook != nil {
+				rulesStore.Hook = tc.ruleStoreHook
+			}
+
+			loki := createTestLokiBackend(t, instrumenttest.NewFakeRequester(), metrics.NewHistorianMetrics(prometheus.NewRegistry(), metrics.Subsystem))
+			loki.ruleStore = rulesStore
+			loki.ac = ac
+
+			// Test conditions.
+			result, err := loki.getFolderUIDsForFilter(context.Background(), models.HistoryQuery{OrgID: orgID, RuleUID: tc.rule.UID, SignedInUser: usr})
+			assert.NoError(t, err)
+			assert.ElementsMatch(t, tc.expectedFolders, result)
+		})
+	}
+}
+
 func createTestLokiBackend(t *testing.T, req alertingInstrument.Requester, met *metrics.Historian) *RemoteLokiBackend {
 	url, _ := url.Parse("http://some.url")
 	cfg := lokiclient.LokiConfig{

pkg/services/ngalert/store/alert_rule.go

@@ -239,6 +239,27 @@ func (st DBstore) GetAlertRuleVersions(ctx context.Context, orgID int64, guid st
 	return alertRules, nil
 }
 
+// GetAlertRuleVersionFolders retrieves a list of unique folder UIDs that the given rule guid has belonged to.
+// Returned slice is ordered with more recent folders first.
+func (st DBstore) GetAlertRuleVersionFolders(ctx context.Context, orgID int64, guid string) ([]string, error) {
+	folders := make([]string, 0)
+	err := st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
+		if err := sess.Table(new(alertRuleVersion)).
+			Select("rule_namespace_uid").
+			Where("rule_org_id = ? AND rule_guid = ?", orgID, guid).
+			GroupBy("rule_namespace_uid").
+			OrderBy("MAX(version) DESC").
+			Find(&folders); err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return folders, nil
+}
+
 // ListDeletedRules retrieves a list of deleted alert rules for the specified organization ID from the database.
 // It ensures that only the latest version of each rule is included and filters out invalid or duplicated versions.
 // Returns a slice of *models.AlertRule  or an error if the operation fails.

pkg/services/ngalert/store/alert_rule_test.go

@@ -1727,6 +1727,55 @@ func TestIntegrationGetRuleVersions(t *testing.T) {
 	})
 }
 
+func TestIntegrationGetAlertRuleVersionFolders(t *testing.T) {
+	tutil.SkipIntegrationTestInShortMode(t)
+
+	// Setup.
+	cfg := setting.NewCfg()
+	cfg.UnifiedAlerting = setting.UnifiedAlertingSettings{BaseInterval: time.Duration(rand.Int64N(100)+1) * time.Second}
+	sqlStore := db.InitTestDB(t)
+	folderService := setupFolderService(t, sqlStore, cfg, featuremgmt.WithFeatures())
+	b := &fakeBus{}
+	store := createTestStore(sqlStore, folderService, &logtest.Fake{}, cfg.UnifiedAlerting, b)
+	orgID := int64(1)
+	gen := models.RuleGen
+	gen = gen.With(gen.WithIntervalMatching(store.Cfg.BaseInterval), gen.WithOrgID(orgID), gen.WithVersion(1))
+
+	inserted, err := store.InsertAlertRules(context.Background(), &models.AlertingUserUID, []models.InsertRule{{AlertRule: gen.Generate()}})
+	require.NoError(t, err)
+	ruleV1, err := store.GetAlertRuleByUID(context.Background(), &models.GetAlertRuleByUIDQuery{UID: inserted[0].UID})
+	require.NoError(t, err)
+
+	oldRule := ruleV1
+	updatedRule := ruleV1
+	updateRule := func(title string, folderUID string) {
+		oldRule = updatedRule
+		updatedRule = models.CopyRule(oldRule, gen.WithTitle(title), gen.WithNamespaceUID(folderUID))
+		require.NoError(t, store.UpdateAlertRules(context.Background(), &models.AlertingUserUID, []models.UpdateRule{{Existing: oldRule, New: *updatedRule}}))
+		updatedRule.Version++ // Simulate version increment after update to avoid conflict errors.
+	}
+
+	// Update rule a couple of times to create versions.
+	originalFolder := oldRule.NamespaceUID
+	updateRule(util.GenerateShortUID(), originalFolder)
+	updateRule(util.GenerateShortUID(), "newfolder-1")
+	updateRule(util.GenerateShortUID(), "newfolder-2")
+	updateRule(util.GenerateShortUID(), "newfolder-2")
+	updateRule(util.GenerateShortUID(), originalFolder)
+	updateRule(util.GenerateShortUID(), "current-folder")
+
+	t.Run("should return rule versions folders sorted in decreasing order", func(t *testing.T) {
+		historicalFolders, err := store.GetAlertRuleVersionFolders(context.Background(), updatedRule.OrgID, updatedRule.GUID)
+		require.NoError(t, err)
+		assert.Equal(t, []string{ // Return folders with more recent first.
+			"current-folder",
+			originalFolder,
+			"newfolder-2",
+			"newfolder-1",
+		}, historicalFolders)
+	})
+}
+
 // createAlertRule creates an alert rule in the database and returns it.
 // If a generator is not specified, uniqueness of primary key is not guaranteed.
 func createRule(tb testing.TB, store *DBstore, generator *models.AlertRuleGenerator) *models.AlertRule {

pkg/services/ngalert/tests/fakes/rules.go

@@ -2,6 +2,7 @@ package fakes
 
 import (
 	"context"
+	"maps"
 	"math/rand"
 	"slices"
 	"strings"
@@ -612,6 +613,30 @@ func (f *RuleStore) GetAlertRuleVersions(_ context.Context, orgID int64, guid st
 	return f.History[guid], nil
 }
 
+func (f *RuleStore) GetAlertRuleVersionFolders(_ context.Context, orgID int64, guid string) ([]string, error) {
+	f.mtx.Lock()
+	defer f.mtx.Unlock()
+
+	q := GenericRecordedQuery{
+		Name:   "GetAlertRuleVersionFolders",
+		Params: []any{orgID, guid},
+	}
+	defer func() {
+		f.RecordedOps = append(f.RecordedOps, q)
+	}()
+
+	if err := f.Hook(q); err != nil {
+		return nil, err
+	}
+
+	folderSet := make(map[string]struct{})
+	for _, rule := range f.History[guid] {
+		folderSet[rule.NamespaceUID] = struct{}{}
+	}
+
+	return slices.Collect(maps.Keys(folderSet)), nil
+}
+
 func (f *RuleStore) ListDeletedRules(_ context.Context, orgID int64) ([]*models.AlertRule, error) {
 	f.mtx.Lock()
 	defer f.mtx.Unlock()