Skip to content

Commit cc97829

Browse files
qinxx108yeya24
authored andcommitted
add api callback
Signed-off-by: Yijie Qin <[email protected]>
1 parent d7b4f0c commit cc97829

File tree

5 files changed

+329
-3
lines changed

5 files changed

+329
-3
lines changed

api/api.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ import (
2020
"runtime"
2121
"time"
2222

23+
"github.com/prometheus/alertmanager/util/callback"
24+
2325
"github.com/go-kit/log"
2426
"github.com/prometheus/client_golang/prometheus"
2527
"github.com/prometheus/common/model"
@@ -74,6 +76,9 @@ type Options struct {
7476
// according to the current active configuration. Alerts returned are
7577
// filtered by the arguments provided to the function.
7678
GroupFunc func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string)
79+
80+
// APICallback define the callback function that each api call will perform before returned.
81+
APICallback callback.Callback
7782
}
7883

7984
func (o Options) validate() error {
@@ -124,6 +129,7 @@ func New(opts Options) (*API, error) {
124129
opts.GroupFunc,
125130
opts.StatusFunc,
126131
opts.Silences,
132+
opts.APICallback,
127133
opts.Peer,
128134
log.With(l, "version", "v2"),
129135
opts.Registry,

api/v2/api.go

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"sync"
2222
"time"
2323

24+
"github.com/prometheus/alertmanager/util/callback"
2425
"github.com/go-kit/log"
2526
"github.com/go-kit/log/level"
2627
"github.com/go-openapi/analysis"
@@ -58,6 +59,7 @@ type API struct {
5859
alerts provider.Alerts
5960
alertGroups groupsFn
6061
getAlertStatus getAlertStatusFn
62+
apiCallback callback.Callback
6163
uptime time.Time
6264

6365
// mtx protects alertmanagerConfig, setAlertStatus and route.
@@ -86,16 +88,21 @@ func NewAPI(
8688
gf groupsFn,
8789
sf getAlertStatusFn,
8890
silences *silence.Silences,
91+
apiCallback callback.Callback,
8992
peer cluster.ClusterPeer,
9093
l log.Logger,
9194
r prometheus.Registerer,
9295
) (*API, error) {
96+
if apiCallback == nil {
97+
apiCallback = callback.NoopAPICallback{}
98+
}
9399
api := API{
94100
alerts: alerts,
95101
getAlertStatus: sf,
96102
alertGroups: gf,
97103
peer: peer,
98104
silences: silences,
105+
apiCallback: apiCallback,
99106
logger: l,
100107
m: metrics.NewAlerts("v2", r),
101108
uptime: time.Now(),
@@ -302,7 +309,13 @@ func (api *API) getAlertsHandler(params alert_ops.GetAlertsParams) middleware.Re
302309
return *res[i].Fingerprint < *res[j].Fingerprint
303310
})
304311

305-
return alert_ops.NewGetAlertsOK().WithPayload(res)
312+
callbackRes, err := api.apiCallback.V2GetAlertsCallback(res)
313+
if err != nil {
314+
level.Error(logger).Log("msg", "Failed to call api callback", "err", err)
315+
return alert_ops.NewGetAlertsInternalServerError().WithPayload(err.Error())
316+
}
317+
318+
return alert_ops.NewGetAlertsOK().WithPayload(callbackRes)
306319
}
307320

308321
func (api *API) postAlertsHandler(params alert_ops.PostAlertsParams) middleware.Responder {
@@ -421,7 +434,13 @@ func (api *API) getAlertGroupsHandler(params alertgroup_ops.GetAlertGroupsParams
421434
res = append(res, ag)
422435
}
423436

424-
return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(res)
437+
callbackRes, err := api.apiCallback.V2GetAlertGroupsCallback(res)
438+
if err != nil {
439+
level.Error(logger).Log("msg", "Failed to call api callback", "err", err)
440+
return alertgroup_ops.NewGetAlertGroupsInternalServerError().WithPayload(err.Error())
441+
}
442+
443+
return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(callbackRes)
425444
}
426445

427446
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: 224 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,18 @@ package v2
1515

1616
import (
1717
"bytes"
18+
"encoding/json"
1819
"fmt"
1920
"io"
2021
"net/http"
2122
"net/http/httptest"
2223
"strconv"
2324
"testing"
2425
"time"
25-
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"
2630
"github.com/go-openapi/runtime"
2731
"github.com/go-openapi/strfmt"
2832
"github.com/prometheus/common/model"
@@ -510,3 +514,222 @@ receivers:
510514
require.Equal(t, tc.body, string(body))
511515
}
512516
}
517+
518+
func TestListAlertsHandler(t *testing.T) {
519+
now := time.Now()
520+
alerts := []*types.Alert{
521+
{
522+
Alert: model.Alert{
523+
Labels: model.LabelSet{"alertname": "alert1"},
524+
StartsAt: now.Add(-time.Minute),
525+
},
526+
},
527+
{
528+
Alert: model.Alert{
529+
Labels: model.LabelSet{"alertname": "alert2"},
530+
StartsAt: now.Add(-time.Minute),
531+
},
532+
},
533+
{
534+
Alert: model.Alert{
535+
Labels: model.LabelSet{"alertname": "alert3"},
536+
StartsAt: now.Add(-time.Minute),
537+
},
538+
},
539+
{
540+
Alert: model.Alert{
541+
Labels: model.LabelSet{"alertname": "alert4"},
542+
StartsAt: now.Add(-time.Minute),
543+
},
544+
},
545+
{
546+
Alert: model.Alert{
547+
Labels: model.LabelSet{"alertname": "alert5"},
548+
StartsAt: now.Add(-time.Minute),
549+
},
550+
},
551+
}
552+
553+
for _, tc := range []struct {
554+
name string
555+
expectedCode int
556+
anames []string
557+
callback callback.Callback
558+
}{
559+
{
560+
"no call back",
561+
200,
562+
[]string{"alert3", "alert2", "alert1", "alert5", "alert4"},
563+
callback.NoopAPICallback{},
564+
},
565+
{
566+
"callback: only return 1 alert",
567+
200,
568+
[]string{"alert3"},
569+
limitNumberOfAlertsReturnedCallback{limit: 1},
570+
},
571+
{
572+
"callback: only return 3 alert",
573+
200,
574+
[]string{"alert3", "alert2", "alert1"},
575+
limitNumberOfAlertsReturnedCallback{limit: 3},
576+
},
577+
} {
578+
t.Run(tc.name, func(t *testing.T) {
579+
alertsProvider := newFakeAlerts(alerts)
580+
api := API{
581+
uptime: time.Now(),
582+
getAlertStatus: getAlertStatus,
583+
logger: log.NewNopLogger(),
584+
apiCallback: tc.callback,
585+
alerts: alertsProvider,
586+
setAlertStatus: func(model.LabelSet) {},
587+
}
588+
api.route = dispatch.NewRoute(&config.Route{Receiver: "def-receiver"}, nil)
589+
r, err := http.NewRequest("GET", "/api/v2/alerts", nil)
590+
require.NoError(t, err)
591+
592+
w := httptest.NewRecorder()
593+
p := runtime.TextProducer()
594+
silence := false
595+
inhibited := false
596+
active := true
597+
responder := api.getAlertsHandler(alert_ops.GetAlertsParams{
598+
HTTPRequest: r,
599+
Silenced: &silence,
600+
Inhibited: &inhibited,
601+
Active: &active,
602+
})
603+
responder.WriteResponse(w, p)
604+
body, _ := io.ReadAll(w.Result().Body)
605+
606+
require.Equal(t, tc.expectedCode, w.Code)
607+
retAlerts := open_api_models.GettableAlerts{}
608+
err = json.Unmarshal(body, &retAlerts)
609+
if err != nil {
610+
t.Fatalf("Unexpected error %v", err)
611+
}
612+
anames := []string{}
613+
for _, a := range retAlerts {
614+
name, ok := a.Labels["alertname"]
615+
if ok {
616+
anames = append(anames, string(name))
617+
}
618+
}
619+
require.Equal(t, tc.anames, anames)
620+
})
621+
}
622+
}
623+
624+
func TestGetAlertGroupsHandler(t *testing.T) {
625+
var startAt time.Time
626+
alerts := []*types.Alert{
627+
{
628+
Alert: model.Alert{
629+
Labels: model.LabelSet{"state": "active", "alertname": "alert1"},
630+
StartsAt: startAt,
631+
},
632+
},
633+
{
634+
Alert: model.Alert{
635+
Labels: model.LabelSet{"state": "unprocessed", "alertname": "alert2"},
636+
StartsAt: startAt,
637+
},
638+
},
639+
}
640+
aginfos := dispatch.AlertGroups{
641+
&dispatch.AlertGroup{
642+
Labels: model.LabelSet{
643+
"alertname": "TestingAlert",
644+
},
645+
Receiver: "testing",
646+
Alerts: alerts[:1],
647+
},
648+
&dispatch.AlertGroup{
649+
Labels: model.LabelSet{
650+
"alertname": "HighErrorRate",
651+
},
652+
Receiver: "prod",
653+
Alerts: alerts[:2],
654+
},
655+
}
656+
for _, tc := range []struct {
657+
name string
658+
numberOfAG int
659+
expectedCode int
660+
callback callback.Callback
661+
}{
662+
{
663+
"no call back",
664+
2,
665+
200,
666+
callback.NoopAPICallback{},
667+
},
668+
{
669+
"callback: only return 1 alert group",
670+
1,
671+
200,
672+
limitNumberOfAlertsReturnedCallback{limit: 1},
673+
},
674+
{
675+
"callback: only return 2 alert group",
676+
2,
677+
200,
678+
limitNumberOfAlertsReturnedCallback{limit: 2},
679+
},
680+
} {
681+
t.Run(tc.name, func(t *testing.T) {
682+
api := API{
683+
uptime: time.Now(),
684+
alertGroups: func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string) {
685+
return aginfos, nil
686+
},
687+
getAlertStatus: getAlertStatus,
688+
logger: log.NewNopLogger(),
689+
apiCallback: tc.callback,
690+
}
691+
r, err := http.NewRequest("GET", "/api/v2/alertgroups", nil)
692+
require.NoError(t, err)
693+
694+
w := httptest.NewRecorder()
695+
p := runtime.TextProducer()
696+
silence := false
697+
inhibited := false
698+
active := true
699+
responder := api.getAlertGroupsHandler(alertgroup_ops.GetAlertGroupsParams{
700+
HTTPRequest: r,
701+
Silenced: &silence,
702+
Inhibited: &inhibited,
703+
Active: &active,
704+
})
705+
responder.WriteResponse(w, p)
706+
body, _ := io.ReadAll(w.Result().Body)
707+
708+
require.Equal(t, tc.expectedCode, w.Code)
709+
retAlertGroups := open_api_models.AlertGroups{}
710+
err = json.Unmarshal(body, &retAlertGroups)
711+
if err != nil {
712+
t.Fatalf("Unexpected error %v", err)
713+
}
714+
require.Equal(t, tc.numberOfAG, len(retAlertGroups))
715+
})
716+
}
717+
}
718+
719+
type limitNumberOfAlertsReturnedCallback struct {
720+
limit int
721+
}
722+
723+
func (n limitNumberOfAlertsReturnedCallback) V2GetAlertsCallback(alerts open_api_models.GettableAlerts) (open_api_models.GettableAlerts, error) {
724+
return alerts[:n.limit], nil
725+
}
726+
727+
func (n limitNumberOfAlertsReturnedCallback) V2GetAlertGroupsCallback(alertgroups open_api_models.AlertGroups) (open_api_models.AlertGroups, error) {
728+
return alertgroups[:n.limit], nil
729+
}
730+
731+
func getAlertStatus(model.Fingerprint) types.AlertStatus {
732+
status := types.AlertStatus{SilencedBy: []string{}, InhibitedBy: []string{}}
733+
status.State = types.AlertStateActive
734+
return status
735+
}

api/v2/testing.go

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

21+
"github.com/prometheus/common/model"
22+
23+
"github.com/prometheus/alertmanager/provider"
24+
"github.com/prometheus/alertmanager/types"
25+
2126
"github.com/go-openapi/strfmt"
2227
"github.com/stretchr/testify/require"
2328

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

0 commit comments

Comments
 (0)