Alerting: Fetch alert rule provenances for a page of rules only (#115643)

* Alerting: Fetch alert rule provenances for a page of rules only

* error when failed to fetch provenance

pkg/services/ngalert/api/api_prometheus_test.go

@@ -2369,6 +2369,140 @@ func TestRouteGetRuleStatuses(t *testing.T) {
 			}
 		})
 
+		t.Run("multi-page pagination loads provenance correctly", func(t *testing.T) {
+			fakeStore, fakeAIM, api, fakeProvisioning := setupAPIFull(t)
+
+			// Create 3 groups with 1 rule each: groups 1 and 3 firing, group 2 normal
+			for i := 1; i <= 3; i++ {
+				rule := gen.With(gen.WithOrgID(orgID), func(r *ngmodels.AlertRule) {
+					r.NamespaceUID = "ns-1"
+					r.RuleGroup = fmt.Sprintf("group-%d", i)
+					r.UID = fmt.Sprintf("rule-%d", i)
+				}, withClassicConditionSingleQuery()).GenerateRef()
+
+				alertState := eval.Normal
+				if i != 2 {
+					alertState = eval.Alerting
+				}
+				fakeAIM.GenerateAlertInstances(orgID, rule.UID, 1, func(s *state.State) *state.State {
+					s.State = alertState
+					s.Labels = data.Labels{"test": "label"}
+					return s
+				})
+				fakeStore.PutRule(context.Background(), rule)
+			}
+
+			// Set provenance for all rules
+			err := fakeProvisioning.SetProvenance(context.Background(),
+				&ngmodels.AlertRule{UID: "rule-1", OrgID: orgID}, orgID, ngmodels.ProvenanceAPI)
+			require.NoError(t, err)
+			err = fakeProvisioning.SetProvenance(context.Background(),
+				&ngmodels.AlertRule{UID: "rule-3", OrgID: orgID}, orgID, ngmodels.ProvenanceFile)
+			require.NoError(t, err)
+
+			// Request firing groups with group_limit=2 - fetches multiple pages, skipping group 2
+			req, err := http.NewRequest("GET", "/api/v1/rules?state=firing&group_limit=2", nil)
+			require.NoError(t, err)
+			c := &contextmodel.ReqContext{
+				Context: &web.Context{Req: req},
+				SignedInUser: &user.SignedInUser{
+					OrgID:       orgID,
+					Permissions: queryPermissions,
+				},
+			}
+
+			resp := api.RouteGetRuleStatuses(c)
+			require.Equal(t, http.StatusOK, resp.Status())
+
+			var res apimodels.RuleResponse
+			require.NoError(t, json.Unmarshal(resp.Body(), &res))
+
+			// Should return 2 firing groups
+			require.Len(t, res.Data.RuleGroups, 2)
+			require.Equal(t, "group-1", res.Data.RuleGroups[0].Name)
+			require.Equal(t, apimodels.Provenance(ngmodels.ProvenanceAPI), res.Data.RuleGroups[0].Rules[0].Provenance)
+			require.Equal(t, "group-3", res.Data.RuleGroups[1].Name)
+			require.Equal(t, apimodels.Provenance(ngmodels.ProvenanceFile), res.Data.RuleGroups[1].Rules[0].Provenance)
+		})
+
+		t.Run("provenance fetch error returns error response in paginated mode", func(t *testing.T) {
+			fakeStore, fakeAIM, api, fakeProvisioning := setupAPIFull(t)
+
+			rule := gen.With(gen.WithOrgID(orgID), func(r *ngmodels.AlertRule) {
+				r.NamespaceUID = "ns-1"
+				r.RuleGroup = "group-1"
+				r.UID = "rule-1"
+			}, withClassicConditionSingleQuery()).GenerateRef()
+
+			fakeAIM.GenerateAlertInstances(orgID, rule.UID, 1, func(s *state.State) *state.State {
+				s.State = eval.Alerting
+				s.Labels = data.Labels{"test": "label"}
+				return s
+			})
+			fakeStore.PutRule(context.Background(), rule)
+
+			fakeProvisioning.GetProvenancesByUIDsFunc = func(ctx context.Context, orgID int64, resourceType string, uids []string) (map[string]ngmodels.Provenance, error) {
+				return nil, errors.New("database connection failed")
+			}
+
+			req, err := http.NewRequest("GET", "/api/v1/rules?group_limit=10", nil)
+			require.NoError(t, err)
+			c := &contextmodel.ReqContext{
+				Context: &web.Context{Req: req},
+				SignedInUser: &user.SignedInUser{
+					OrgID:       orgID,
+					Permissions: queryPermissions,
+				},
+			}
+
+			resp := api.RouteGetRuleStatuses(c)
+			require.Equal(t, http.StatusInternalServerError, resp.Status())
+
+			var res apimodels.RuleResponse
+			require.NoError(t, json.Unmarshal(resp.Body(), &res))
+			require.Equal(t, "error", res.Status)
+			require.Contains(t, res.Error, "failed to load provenance")
+		})
+
+		t.Run("provenance fetch error returns error response in non-paginated mode", func(t *testing.T) {
+			fakeStore, fakeAIM, api, fakeProvisioning := setupAPIFull(t)
+
+			rule := gen.With(gen.WithOrgID(orgID), func(r *ngmodels.AlertRule) {
+				r.NamespaceUID = "ns-1"
+				r.RuleGroup = "group-1"
+				r.UID = "rule-1"
+			}, withClassicConditionSingleQuery()).GenerateRef()
+
+			fakeAIM.GenerateAlertInstances(orgID, rule.UID, 1, func(s *state.State) *state.State {
+				s.State = eval.Alerting
+				s.Labels = data.Labels{"test": "label"}
+				return s
+			})
+			fakeStore.PutRule(context.Background(), rule)
+
+			fakeProvisioning.GetProvenancesFunc = func(ctx context.Context, orgID int64, resourceType string) (map[string]ngmodels.Provenance, error) {
+				return nil, errors.New("database connection failed")
+			}
+
+			req, err := http.NewRequest("GET", "/api/v1/rules", nil)
+			require.NoError(t, err)
+			c := &contextmodel.ReqContext{
+				Context: &web.Context{Req: req},
+				SignedInUser: &user.SignedInUser{
+					OrgID:       orgID,
+					Permissions: queryPermissions,
+				},
+			}
+
+			resp := api.RouteGetRuleStatuses(c)
+			require.Equal(t, http.StatusInternalServerError, resp.Status())
+
+			var res apimodels.RuleResponse
+			require.NoError(t, json.Unmarshal(resp.Body(), &res))
+			require.Equal(t, "error", res.Status)
+			require.Contains(t, res.Error, "failed to load provenance")
+		})
+
 		t.Run("state filter continues when first page has no matches", func(t *testing.T) {
 			fakeStore, fakeAIM, api := setupAPI(t)
 

pkg/services/ngalert/api/prometheus/api_prometheus.go

@@ -54,6 +54,7 @@ type StatusReader interface {
 
 type ProvenanceStore interface {
 	GetProvenances(ctx context.Context, org int64, resourceType string) (map[string]ngmodels.Provenance, error)
+	GetProvenancesByUIDs(ctx context.Context, org int64, resourceType string, uids []string) (map[string]ngmodels.Provenance, error)
 }
 
 type PrometheusSrv struct {
@@ -328,14 +329,6 @@ func (srv PrometheusSrv) RouteGetRuleStatuses(c *contextmodel.ReqContext) respon
 	span.AddEvent("User permissions checked")
 	span.SetAttributes(attribute.Int("allowedNamespaces", len(allowedNamespaces)))
 
-	provenanceRecords, err := srv.provenanceStore.GetProvenances(c.Req.Context(), c.GetOrgID(), (&ngmodels.AlertRule{}).ResourceType())
-	if err != nil {
-		ruleResponse.Status = "error"
-		ruleResponse.Error = fmt.Sprintf("failed to get provenances visible to the user: %s", err.Error())
-		ruleResponse.ErrorType = apiv1.ErrServer
-		return response.JSON(ruleResponse.HTTPStatusCode(), ruleResponse)
-	}
-
 	ruleResponse = PrepareRuleGroupStatusesV2(
 		srv.log,
 		srv.store,
@@ -347,7 +340,7 @@ func (srv PrometheusSrv) RouteGetRuleStatuses(c *contextmodel.ReqContext) respon
 		},
 		RuleStatusMutatorGenerator(srv.status),
 		RuleAlertStateMutatorGenerator(srv.manager),
-		provenanceRecords,
+		srv.provenanceStore,
 	)
 
 	return response.JSON(ruleResponse.HTTPStatusCode(), ruleResponse)
@@ -454,6 +447,7 @@ func RuleAlertStateMutatorGenerator(manager state.AlertInstanceManager) RuleAler
 type paginationContext struct {
 	opts              RuleGroupStatusesOptions
 	provenanceRecords map[string]ngmodels.Provenance
+	provenanceStore   ProvenanceStore
 	ruleStatusMutator RuleStatusMutator
 	alertStateMutator RuleAlertStateMutator
 
@@ -532,6 +526,37 @@ func (ctx *paginationContext) fetchAndFilterPage(log log.Logger, store ListAlert
 	)
 	span.AddEvent("Alert rules retrieved from store")
 
+	// Load provenance for this page's rules
+	if ctx.provenanceStore != nil {
+		maxGroups := getInt64WithDefault(ctx.opts.Query, "group_limit", -1)
+		maxRules := getInt64WithDefault(ctx.opts.Query, "rule_limit", -1)
+
+		if maxGroups > 0 || maxRules > 0 {
+			// Paginated, fetch and merge provenances for this page
+			uids := make([]string, 0, len(ruleList))
+			for _, rule := range ruleList {
+				uids = append(uids, rule.UID)
+			}
+			pageProvenances, err := ctx.provenanceStore.GetProvenancesByUIDs(ctx.opts.Ctx, ctx.opts.OrgID, (&ngmodels.AlertRule{}).ResourceType(), uids)
+			if err != nil {
+				return pageResult{}, fmt.Errorf("failed to load provenance: %w", err)
+			}
+			if ctx.provenanceRecords == nil {
+				ctx.provenanceRecords = pageProvenances
+			} else {
+				maps.Copy(ctx.provenanceRecords, pageProvenances)
+			}
+		} else if ctx.provenanceRecords == nil {
+			// Not paginated, fetch all once
+			var err error
+			ctx.provenanceRecords, err = ctx.provenanceStore.GetProvenances(ctx.opts.Ctx, ctx.opts.OrgID, (&ngmodels.AlertRule{}).ResourceType())
+			if err != nil {
+				return pageResult{}, fmt.Errorf("failed to load provenance: %w", err)
+			}
+		}
+	}
+	span.AddEvent("Provenances retrieved from store")
+
 	groupedRules := getGroupedRules(log, ruleList, ctx.ruleNamesSet, ctx.opts.AllowedNamespaces)
 
 	result := pageResult{
@@ -643,7 +668,7 @@ func paginateRuleGroups(log log.Logger, store ListAlertRulesStoreV2, ctx *pagina
 	return allGroups, rulesTotals, continueToken, nil
 }
 
-func PrepareRuleGroupStatusesV2(log log.Logger, store ListAlertRulesStoreV2, opts RuleGroupStatusesOptions, ruleStatusMutator RuleStatusMutator, alertStateMutator RuleAlertStateMutator, provenanceRecords map[string]ngmodels.Provenance) apimodels.RuleResponse {
+func PrepareRuleGroupStatusesV2(log log.Logger, store ListAlertRulesStoreV2, opts RuleGroupStatusesOptions, ruleStatusMutator RuleStatusMutator, alertStateMutator RuleAlertStateMutator, provenanceStore ProvenanceStore) apimodels.RuleResponse {
 	ctx, span := tracer.Start(opts.Ctx, "api.prometheus.PrepareRuleGroupStatusesV2")
 	defer span.End()
 	opts.Ctx = ctx
@@ -835,7 +860,8 @@ func PrepareRuleGroupStatusesV2(log log.Logger, store ListAlertRulesStoreV2, opt
 	span.SetAttributes(attribute.Bool("compact", compact))
 	pagCtx := &paginationContext{
 		opts:               opts,
-		provenanceRecords:  provenanceRecords,
+		provenanceRecords:  nil,
+		provenanceStore:    provenanceStore,
 		ruleStatusMutator:  ruleStatusMutator,
 		alertStateMutator:  alertStateMutator,
 		namespaceUIDs:      namespaceUIDs,

pkg/services/ngalert/notifier/alertmanager_config.go

@@ -485,6 +485,7 @@ func assignReceiverConfigsUIDs(c []*definitions.PostableApiReceiver) error {
 type provisioningStore interface {
 	GetProvenance(ctx context.Context, o models.Provisionable, org int64) (models.Provenance, error)
 	GetProvenances(ctx context.Context, org int64, resourceType string) (map[string]models.Provenance, error)
+	GetProvenancesByUIDs(ctx context.Context, org int64, resourceType string, uids []string) (map[string]models.Provenance, error)
 	SetProvenance(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error
 	DeleteProvenance(ctx context.Context, o models.Provisionable, org int64) error
 }

pkg/services/ngalert/provisioning/persist.go

@@ -19,6 +19,7 @@ type alertmanagerConfigStore interface {
 type ProvisioningStore interface {
 	GetProvenance(ctx context.Context, o models.Provisionable, org int64) (models.Provenance, error)
 	GetProvenances(ctx context.Context, org int64, resourceType string) (map[string]models.Provenance, error)
+	GetProvenancesByUIDs(ctx context.Context, org int64, resourceType string, uids []string) (map[string]models.Provenance, error)
 	SetProvenance(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error
 	DeleteProvenance(ctx context.Context, o models.Provisionable, org int64) error
 }

pkg/services/ngalert/provisioning/provisioning_store_mock.go

@@ -188,6 +188,67 @@ func (_c *MockProvisioningStore_GetProvenances_Call) RunAndReturn(run func(conte
 	return _c
 }
 
+// GetProvenancesByUIDs provides a mock function with given fields: ctx, org, resourceType, uids
+func (_m *MockProvisioningStore) GetProvenancesByUIDs(ctx context.Context, org int64, resourceType string, uids []string) (map[string]models.Provenance, error) {
+	ret := _m.Called(ctx, org, resourceType, uids)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetProvenancesByUIDs")
+	}
+
+	var r0 map[string]models.Provenance
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, int64, string, []string) (map[string]models.Provenance, error)); ok {
+		return rf(ctx, org, resourceType, uids)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, int64, string, []string) map[string]models.Provenance); ok {
+		r0 = rf(ctx, org, resourceType, uids)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(map[string]models.Provenance)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, int64, string, []string) error); ok {
+		r1 = rf(ctx, org, resourceType, uids)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockProvisioningStore_GetProvenancesByUIDs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetProvenancesByUIDs'
+type MockProvisioningStore_GetProvenancesByUIDs_Call struct {
+	*mock.Call
+}
+
+// GetProvenancesByUIDs is a helper method to define mock.On call
+//   - ctx context.Context
+//   - org int64
+//   - resourceType string
+//   - uids []string
+func (_e *MockProvisioningStore_Expecter) GetProvenancesByUIDs(ctx interface{}, org interface{}, resourceType interface{}, uids interface{}) *MockProvisioningStore_GetProvenancesByUIDs_Call {
+	return &MockProvisioningStore_GetProvenancesByUIDs_Call{Call: _e.mock.On("GetProvenancesByUIDs", ctx, org, resourceType, uids)}
+}
+
+func (_c *MockProvisioningStore_GetProvenancesByUIDs_Call) Run(run func(ctx context.Context, org int64, resourceType string, uids []string)) *MockProvisioningStore_GetProvenancesByUIDs_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(int64), args[2].(string), args[3].([]string))
+	})
+	return _c
+}
+
+func (_c *MockProvisioningStore_GetProvenancesByUIDs_Call) Return(_a0 map[string]models.Provenance, _a1 error) *MockProvisioningStore_GetProvenancesByUIDs_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockProvisioningStore_GetProvenancesByUIDs_Call) RunAndReturn(run func(context.Context, int64, string, []string) (map[string]models.Provenance, error)) *MockProvisioningStore_GetProvenancesByUIDs_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // SetProvenance provides a mock function with given fields: ctx, o, org, p
 func (_m *MockProvisioningStore) SetProvenance(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error {
 	ret := _m.Called(ctx, o, org, p)

pkg/services/ngalert/store/provisioning_store.go

@@ -62,6 +62,30 @@ func (st DBstore) GetProvenances(ctx context.Context, org int64, resourceType st
 	return resultMap, err
 }
 
+// GetProvenancesByUIDs gets the provenance status for specific UIDs.
+func (st DBstore) GetProvenancesByUIDs(ctx context.Context, org int64, resourceType string, uids []string) (map[string]models.Provenance, error) {
+	if len(uids) == 0 {
+		return map[string]models.Provenance{}, nil
+	}
+
+	result := make(map[string]models.Provenance, len(uids))
+	err := st.SQLStore.WithDbSession(ctx, func(sess *db.Session) error {
+		rawData, err := sess.Table(provenanceRecord{}).
+			Where("record_type = ? AND org_id = ?", resourceType, org).
+			In("record_key", uids).
+			Cols("record_key", "provenance").
+			QueryString()
+		if err != nil {
+			return fmt.Errorf("failed to query for existing provenance status: %w", err)
+		}
+		for _, data := range rawData {
+			result[data["record_key"]] = models.Provenance(data["provenance"])
+		}
+		return nil
+	})
+	return result, err
+}
+
 // SetProvenance changes the provenance status for a provisionable object.
 func (st DBstore) SetProvenance(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error {
 	recordType := o.ResourceType()

pkg/services/ngalert/store/provisioning_store_test.go

@@ -133,6 +133,55 @@ func TestIntegrationProvisioningStore(t *testing.T) {
 				require.Equal(t, models.ProvenanceAPI, p[rule2.UID])
 			})
 
+			t.Run("Store should return provenances by UIDs", func(t *testing.T) {
+				const orgID = 124
+				rule1 := models.AlertRule{UID: "uid-1", OrgID: orgID}
+				rule2 := models.AlertRule{UID: "uid-2", OrgID: orgID}
+				rule3 := models.AlertRule{UID: "uid-3", OrgID: orgID}
+
+				err := store.SetProvenance(context.Background(), &rule1, orgID, models.ProvenanceFile)
+				require.NoError(t, err)
+				err = store.SetProvenance(context.Background(), &rule2, orgID, models.ProvenanceAPI)
+				require.NoError(t, err)
+				err = store.SetProvenance(context.Background(), &rule3, orgID, models.ProvenanceFile)
+				require.NoError(t, err)
+
+				// Fetch only rule1 and rule2
+				p, err := store.GetProvenancesByUIDs(context.Background(), orgID, rule1.ResourceType(), []string{rule1.UID, rule2.UID})
+				require.NoError(t, err)
+				require.Len(t, p, 2)
+				require.Equal(t, models.ProvenanceFile, p[rule1.UID])
+				require.Equal(t, models.ProvenanceAPI, p[rule2.UID])
+				_, exists := p[rule3.UID]
+				require.False(t, exists)
+			})
+
+			t.Run("GetProvenancesByUIDs returns empty map for empty UIDs", func(t *testing.T) {
+				p, err := store.GetProvenancesByUIDs(context.Background(), 1, "alertRule", []string{})
+				require.NoError(t, err)
+				require.Empty(t, p)
+			})
+
+			t.Run("GetProvenancesByUIDs respects org ID", func(t *testing.T) {
+				const orgID1 = 125
+				const orgID2 = 126
+				rule := models.AlertRule{UID: "cross-org-uid"}
+
+				err := store.SetProvenance(context.Background(), &rule, orgID1, models.ProvenanceFile)
+				require.NoError(t, err)
+
+				// Should not find in different org
+				p, err := store.GetProvenancesByUIDs(context.Background(), orgID2, rule.ResourceType(), []string{rule.UID})
+				require.NoError(t, err)
+				require.Empty(t, p)
+
+				// Should find in correct org
+				p, err = store.GetProvenancesByUIDs(context.Background(), orgID1, rule.ResourceType(), []string{rule.UID})
+				require.NoError(t, err)
+				require.Len(t, p, 1)
+				require.Equal(t, models.ProvenanceFile, p[rule.UID])
+			})
+
 			t.Run("Store should delete provenance correctly", func(t *testing.T) {
 				const orgID = 1234
 				ruleOrg := models.AlertRule{

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

@@ -8,12 +8,13 @@ import (
 )
 
 type FakeProvisioningStore struct {
-	Calls                []Call
-	Records              map[int64]map[string]models.Provenance
-	GetProvenanceFunc    func(ctx context.Context, o models.Provisionable, org int64) (models.Provenance, error)
-	GetProvenancesFunc   func(ctx context.Context, orgID int64, resourceType string) (map[string]models.Provenance, error)
-	SetProvenanceFunc    func(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error
-	DeleteProvenanceFunc func(ctx context.Context, o models.Provisionable, org int64) error
+	Calls                    []Call
+	Records                  map[int64]map[string]models.Provenance
+	GetProvenanceFunc        func(ctx context.Context, o models.Provisionable, org int64) (models.Provenance, error)
+	GetProvenancesFunc       func(ctx context.Context, orgID int64, resourceType string) (map[string]models.Provenance, error)
+	GetProvenancesByUIDsFunc func(ctx context.Context, orgID int64, resourceType string, uids []string) (map[string]models.Provenance, error)
+	SetProvenanceFunc        func(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error
+	DeleteProvenanceFunc     func(ctx context.Context, o models.Provisionable, org int64) error
 }
 
 func NewFakeProvisioningStore() *FakeProvisioningStore {
@@ -51,6 +52,23 @@ func (f *FakeProvisioningStore) GetProvenances(ctx context.Context, orgID int64,
 	return results, nil
 }
 
+func (f *FakeProvisioningStore) GetProvenancesByUIDs(ctx context.Context, orgID int64, resourceType string, uids []string) (map[string]models.Provenance, error) {
+	f.Calls = append(f.Calls, Call{MethodName: "GetProvenancesByUIDs", Arguments: []any{ctx, orgID, resourceType, uids}})
+	if f.GetProvenancesByUIDsFunc != nil {
+		return f.GetProvenancesByUIDsFunc(ctx, orgID, resourceType, uids)
+	}
+	results := make(map[string]models.Provenance)
+	if val, ok := f.Records[orgID]; ok {
+		for _, uid := range uids {
+			key := uid + resourceType
+			if prov, ok := val[key]; ok {
+				results[uid] = prov
+			}
+		}
+	}
+	return results, nil
+}
+
 func (f *FakeProvisioningStore) SetProvenance(ctx context.Context, o models.Provisionable, org int64, p models.Provenance) error {
 	f.Calls = append(f.Calls, Call{MethodName: "SetProvenance", Arguments: []any{ctx, o, org, p}})
 	if f.SetProvenanceFunc != nil {