Skip to content

Commit a23237f

Browse files
qinxx108yeya24
qinxx108
authored and
yeya24
committed
add api callback
Signed-off-by: Yijie Qin <[email protected]>
1 parent b2099ea commit a23237f

File tree

5 files changed

+327
-3
lines changed

5 files changed

+327
-3
lines changed

‎api/api.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333
"github.com/prometheus/alertmanager/provider"
3434
"github.com/prometheus/alertmanager/silence"
3535
"github.com/prometheus/alertmanager/types"
36+
"github.com/prometheus/alertmanager/util/callback"
3637
)
3738

3839
// API represents all APIs of Alertmanager.
@@ -79,6 +80,9 @@ type Options struct {
7980
// according to the current active configuration. Alerts returned are
8081
// filtered by the arguments provided to the function.
8182
GroupFunc func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string)
83+
84+
// APICallback define the callback function that each api call will perform before returned.
85+
APICallback callback.Callback
8286
}
8387

8488
func (o Options) validate() error {
@@ -124,6 +128,7 @@ func New(opts Options) (*API, error) {
124128
opts.AlertStatusFunc,
125129
opts.GroupMutedFunc,
126130
opts.Silences,
131+
opts.APICallback,
127132
opts.Peer,
128133
l.With("version", "v2"),
129134
opts.Registry,

‎api/v2/api.go‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ import (
5050
"github.com/prometheus/alertmanager/silence"
5151
"github.com/prometheus/alertmanager/silence/silencepb"
5252
"github.com/prometheus/alertmanager/types"
53+
"github.com/prometheus/alertmanager/util/callback"
5354
)
5455

5556
// API represents an Alertmanager API v2.
@@ -60,6 +61,7 @@ type API struct {
6061
alertGroups groupsFn
6162
getAlertStatus getAlertStatusFn
6263
groupMutedFunc groupMutedFunc
64+
apiCallback callback.Callback
6365
uptime time.Time
6466

6567
// mtx protects alertmanagerConfig, setAlertStatus and route.
@@ -90,17 +92,22 @@ func NewAPI(
9092
asf getAlertStatusFn,
9193
gmf groupMutedFunc,
9294
silences *silence.Silences,
95+
apiCallback callback.Callback,
9396
peer cluster.ClusterPeer,
9497
l *slog.Logger,
9598
r prometheus.Registerer,
9699
) (*API, error) {
100+
if apiCallback == nil {
101+
apiCallback = callback.NoopAPICallback{}
102+
}
97103
api := API{
98104
alerts: alerts,
99105
getAlertStatus: asf,
100106
alertGroups: gf,
101107
groupMutedFunc: gmf,
102108
peer: peer,
103109
silences: silences,
110+
apiCallback: apiCallback,
104111
logger: l,
105112
m: metrics.NewAlerts(r),
106113
uptime: time.Now(),
@@ -307,7 +314,13 @@ func (api *API) getAlertsHandler(params alert_ops.GetAlertsParams) middleware.Re
307314
return *res[i].Fingerprint < *res[j].Fingerprint
308315
})
309316

310-
return alert_ops.NewGetAlertsOK().WithPayload(res)
317+
callbackRes, err := api.apiCallback.V2GetAlertsCallback(res)
318+
if err != nil {
319+
logger.Error("Failed to call api callback", "err", err)
320+
return alert_ops.NewGetAlertsInternalServerError().WithPayload(err.Error())
321+
}
322+
323+
return alert_ops.NewGetAlertsOK().WithPayload(callbackRes)
311324
}
312325

313326
func (api *API) postAlertsHandler(params alert_ops.PostAlertsParams) middleware.Responder {
@@ -431,7 +444,13 @@ func (api *API) getAlertGroupsHandler(params alertgroup_ops.GetAlertGroupsParams
431444
res = append(res, ag)
432445
}
433446

434-
return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(res)
447+
callbackRes, err := api.apiCallback.V2GetAlertGroupsCallback(res)
448+
if err != nil {
449+
logger.Error("Failed to call api callback", "err", err)
450+
return alertgroup_ops.NewGetAlertGroupsInternalServerError().WithPayload(err.Error())
451+
}
452+
453+
return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(callbackRes)
435454
}
436455

437456
func (api *API) alertFilter(matchers []*labels.Matcher, silenced, inhibited, active bool) func(a *types.Alert, now time.Time) bool {

‎api/v2/api_test.go‎

Lines changed: 223 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,10 @@ import (
2323
"strconv"
2424
"testing"
2525
"time"
26-
26+
"github.com/prometheus/alertmanager/dispatch"
27+
alert_ops "github.com/prometheus/alertmanager/api/v2/restapi/operations/alert"
28+
alertgroup_ops "github.com/prometheus/alertmanager/api/v2/restapi/operations/alertgroup"
29+
"github.com/prometheus/alertmanager/util/callback"
2730
"github.com/go-openapi/runtime"
2831
"github.com/go-openapi/runtime/middleware"
2932
"github.com/go-openapi/strfmt"
@@ -584,3 +587,222 @@ receivers:
584587
require.Equal(t, tc.body, string(body))
585588
}
586589
}
590+
591+
func TestListAlertsHandler(t *testing.T) {
592+
now := time.Now()
593+
alerts := []*types.Alert{
594+
{
595+
Alert: model.Alert{
596+
Labels: model.LabelSet{"alertname": "alert1"},
597+
StartsAt: now.Add(-time.Minute),
598+
},
599+
},
600+
{
601+
Alert: model.Alert{
602+
Labels: model.LabelSet{"alertname": "alert2"},
603+
StartsAt: now.Add(-time.Minute),
604+
},
605+
},
606+
{
607+
Alert: model.Alert{
608+
Labels: model.LabelSet{"alertname": "alert3"},
609+
StartsAt: now.Add(-time.Minute),
610+
},
611+
},
612+
{
613+
Alert: model.Alert{
614+
Labels: model.LabelSet{"alertname": "alert4"},
615+
StartsAt: now.Add(-time.Minute),
616+
},
617+
},
618+
{
619+
Alert: model.Alert{
620+
Labels: model.LabelSet{"alertname": "alert5"},
621+
StartsAt: now.Add(-time.Minute),
622+
},
623+
},
624+
}
625+
626+
for _, tc := range []struct {
627+
name string
628+
expectedCode int
629+
anames []string
630+
callback callback.Callback
631+
}{
632+
{
633+
"no call back",
634+
200,
635+
[]string{"alert3", "alert2", "alert1", "alert5", "alert4"},
636+
callback.NoopAPICallback{},
637+
},
638+
{
639+
"callback: only return 1 alert",
640+
200,
641+
[]string{"alert3"},
642+
limitNumberOfAlertsReturnedCallback{limit: 1},
643+
},
644+
{
645+
"callback: only return 3 alert",
646+
200,
647+
[]string{"alert3", "alert2", "alert1"},
648+
limitNumberOfAlertsReturnedCallback{limit: 3},
649+
},
650+
} {
651+
t.Run(tc.name, func(t *testing.T) {
652+
alertsProvider := newFakeAlerts(alerts)
653+
api := API{
654+
uptime: time.Now(),
655+
getAlertStatus: getAlertStatus,
656+
logger: log.NewNopLogger(),
657+
apiCallback: tc.callback,
658+
alerts: alertsProvider,
659+
setAlertStatus: func(model.LabelSet) {},
660+
}
661+
api.route = dispatch.NewRoute(&config.Route{Receiver: "def-receiver"}, nil)
662+
r, err := http.NewRequest("GET", "/api/v2/alerts", nil)
663+
require.NoError(t, err)
664+
665+
w := httptest.NewRecorder()
666+
p := runtime.TextProducer()
667+
silence := false
668+
inhibited := false
669+
active := true
670+
responder := api.getAlertsHandler(alert_ops.GetAlertsParams{
671+
HTTPRequest: r,
672+
Silenced: &silence,
673+
Inhibited: &inhibited,
674+
Active: &active,
675+
})
676+
responder.WriteResponse(w, p)
677+
body, _ := io.ReadAll(w.Result().Body)
678+
679+
require.Equal(t, tc.expectedCode, w.Code)
680+
retAlerts := open_api_models.GettableAlerts{}
681+
err = json.Unmarshal(body, &retAlerts)
682+
if err != nil {
683+
t.Fatalf("Unexpected error %v", err)
684+
}
685+
anames := []string{}
686+
for _, a := range retAlerts {
687+
name, ok := a.Labels["alertname"]
688+
if ok {
689+
anames = append(anames, string(name))
690+
}
691+
}
692+
require.Equal(t, tc.anames, anames)
693+
})
694+
}
695+
}
696+
697+
func TestGetAlertGroupsHandler(t *testing.T) {
698+
var startAt time.Time
699+
alerts := []*types.Alert{
700+
{
701+
Alert: model.Alert{
702+
Labels: model.LabelSet{"state": "active", "alertname": "alert1"},
703+
StartsAt: startAt,
704+
},
705+
},
706+
{
707+
Alert: model.Alert{
708+
Labels: model.LabelSet{"state": "unprocessed", "alertname": "alert2"},
709+
StartsAt: startAt,
710+
},
711+
},
712+
}
713+
aginfos := dispatch.AlertGroups{
714+
&dispatch.AlertGroup{
715+
Labels: model.LabelSet{
716+
"alertname": "TestingAlert",
717+
},
718+
Receiver: "testing",
719+
Alerts: alerts[:1],
720+
},
721+
&dispatch.AlertGroup{
722+
Labels: model.LabelSet{
723+
"alertname": "HighErrorRate",
724+
},
725+
Receiver: "prod",
726+
Alerts: alerts[:2],
727+
},
728+
}
729+
for _, tc := range []struct {
730+
name string
731+
numberOfAG int
732+
expectedCode int
733+
callback callback.Callback
734+
}{
735+
{
736+
"no call back",
737+
2,
738+
200,
739+
callback.NoopAPICallback{},
740+
},
741+
{
742+
"callback: only return 1 alert group",
743+
1,
744+
200,
745+
limitNumberOfAlertsReturnedCallback{limit: 1},
746+
},
747+
{
748+
"callback: only return 2 alert group",
749+
2,
750+
200,
751+
limitNumberOfAlertsReturnedCallback{limit: 2},
752+
},
753+
} {
754+
t.Run(tc.name, func(t *testing.T) {
755+
api := API{
756+
uptime: time.Now(),
757+
alertGroups: func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string) {
758+
return aginfos, nil
759+
},
760+
getAlertStatus: getAlertStatus,
761+
logger: log.NewNopLogger(),
762+
apiCallback: tc.callback,
763+
}
764+
r, err := http.NewRequest("GET", "/api/v2/alertgroups", nil)
765+
require.NoError(t, err)
766+
767+
w := httptest.NewRecorder()
768+
p := runtime.TextProducer()
769+
silence := false
770+
inhibited := false
771+
active := true
772+
responder := api.getAlertGroupsHandler(alertgroup_ops.GetAlertGroupsParams{
773+
HTTPRequest: r,
774+
Silenced: &silence,
775+
Inhibited: &inhibited,
776+
Active: &active,
777+
})
778+
responder.WriteResponse(w, p)
779+
body, _ := io.ReadAll(w.Result().Body)
780+
781+
require.Equal(t, tc.expectedCode, w.Code)
782+
retAlertGroups := open_api_models.AlertGroups{}
783+
err = json.Unmarshal(body, &retAlertGroups)
784+
if err != nil {
785+
t.Fatalf("Unexpected error %v", err)
786+
}
787+
require.Equal(t, tc.numberOfAG, len(retAlertGroups))
788+
})
789+
}
790+
}
791+
792+
type limitNumberOfAlertsReturnedCallback struct {
793+
limit int
794+
}
795+
796+
func (n limitNumberOfAlertsReturnedCallback) V2GetAlertsCallback(alerts open_api_models.GettableAlerts) (open_api_models.GettableAlerts, error) {
797+
return alerts[:n.limit], nil
798+
}
799+
800+
func (n limitNumberOfAlertsReturnedCallback) V2GetAlertGroupsCallback(alertgroups open_api_models.AlertGroups) (open_api_models.AlertGroups, error) {
801+
return alertgroups[:n.limit], nil
802+
}
803+
804+
func getAlertStatus(model.Fingerprint) types.AlertStatus {
805+
status := types.AlertStatus{SilencedBy: []string{}, InhibitedBy: []string{}}
806+
status.State = types.AlertStateActive
807+
return status
808+
}

‎api/v2/testing.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,11 @@ import (
1717
"testing"
1818
"time"
1919

20+
"github.com/prometheus/common/model"
21+
22+
"github.com/prometheus/alertmanager/provider"
23+
"github.com/prometheus/alertmanager/types"
24+
2025
"github.com/go-openapi/strfmt"
2126

2227
open_api_models "github.com/prometheus/alertmanager/api/v2/models"
@@ -63,3 +68,40 @@ func createLabelMatcher(t *testing.T, name, value string, matchType labels.Match
6368
matcher, _ := labels.NewMatcher(matchType, name, value)
6469
return matcher
6570
}
71+
72+
// fakeAlerts is a struct implementing the provider.Alerts interface for tests.
73+
type fakeAlerts struct {
74+
fps map[model.Fingerprint]int
75+
alerts []*types.Alert
76+
err error
77+
}
78+
79+
func newFakeAlerts(alerts []*types.Alert) *fakeAlerts {
80+
fps := make(map[model.Fingerprint]int)
81+
for i, a := range alerts {
82+
fps[a.Fingerprint()] = i
83+
}
84+
f := &fakeAlerts{
85+
alerts: alerts,
86+
fps: fps,
87+
}
88+
return f
89+
}
90+
91+
func (f *fakeAlerts) Subscribe() provider.AlertIterator { return nil }
92+
func (f *fakeAlerts) Get(model.Fingerprint) (*types.Alert, error) { return nil, nil }
93+
func (f *fakeAlerts) Put(alerts ...*types.Alert) error {
94+
return f.err
95+
}
96+
97+
func (f *fakeAlerts) GetPending() provider.AlertIterator {
98+
ch := make(chan *types.Alert)
99+
done := make(chan struct{})
100+
go func() {
101+
defer close(ch)
102+
for _, a := range f.alerts {
103+
ch <- a
104+
}
105+
}()
106+
return provider.NewAlertIterator(ch, done, f.err)
107+
}

0 commit comments

Comments
 (0)