diff --git a/pkg/services/ngalert/api/api_prometheus_test.go b/pkg/services/ngalert/api/api_prometheus_test.go index 75ec6c901fd..dc0f6d1f13d 100644 --- a/pkg/services/ngalert/api/api_prometheus_test.go +++ b/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) diff --git a/pkg/services/ngalert/api/prometheus/api_prometheus.go b/pkg/services/ngalert/api/prometheus/api_prometheus.go index b4e14a66cfe..077761d0caf 100644 --- a/pkg/services/ngalert/api/prometheus/api_prometheus.go +++ b/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, diff --git a/pkg/services/ngalert/notifier/alertmanager_config.go b/pkg/services/ngalert/notifier/alertmanager_config.go index 7e755903791..ba8d37809b1 100644 --- a/pkg/services/ngalert/notifier/alertmanager_config.go +++ b/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 } diff --git a/pkg/services/ngalert/provisioning/persist.go b/pkg/services/ngalert/provisioning/persist.go index 914a3644984..e6d6b37fc24 100644 --- a/pkg/services/ngalert/provisioning/persist.go +++ b/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 } diff --git a/pkg/services/ngalert/provisioning/provisioning_store_mock.go b/pkg/services/ngalert/provisioning/provisioning_store_mock.go index 31cc77e26a6..bbc115d87c1 100644 --- a/pkg/services/ngalert/provisioning/provisioning_store_mock.go +++ b/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) diff --git a/pkg/services/ngalert/store/provisioning_store.go b/pkg/services/ngalert/store/provisioning_store.go index 27f03143c98..df5d7dc80c6 100644 --- a/pkg/services/ngalert/store/provisioning_store.go +++ b/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() diff --git a/pkg/services/ngalert/store/provisioning_store_test.go b/pkg/services/ngalert/store/provisioning_store_test.go index b6f8b8fe5cd..359d594da0a 100644 --- a/pkg/services/ngalert/store/provisioning_store_test.go +++ b/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{ diff --git a/pkg/services/ngalert/tests/fakes/provisioning.go b/pkg/services/ngalert/tests/fakes/provisioning.go index 43de0a6dc68..fce2586f120 100644 --- a/pkg/services/ngalert/tests/fakes/provisioning.go +++ b/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 {