add api callback

Signed-off-by: Yijie Qin <[email protected]>

Author: qinxx108

api/api.go

@@ -20,6 +20,8 @@ import (
 	"runtime"
 	"time"
 
+	"github.com/prometheus/alertmanager/util/callback"
+
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
@@ -74,6 +76,9 @@ type Options struct {
 	// according to the current active configuration. Alerts returned are
 	// filtered by the arguments provided to the function.
 	GroupFunc func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string)
+
+	// APICallback define the callback function that each api call will perform before returned.
+	APICallback callback.Callback
 }
 
 func (o Options) validate() error {
@@ -124,6 +129,7 @@ func New(opts Options) (*API, error) {
 		opts.GroupFunc,
 		opts.StatusFunc,
 		opts.Silences,
+		opts.APICallback,
 		opts.Peer,
 		log.With(l, "version", "v2"),
 		opts.Registry,

api/v2/api.go

@@ -21,6 +21,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/prometheus/alertmanager/util/callback"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/go-openapi/analysis"
@@ -58,6 +59,7 @@ type API struct {
 	alerts         provider.Alerts
 	alertGroups    groupsFn
 	getAlertStatus getAlertStatusFn
+	apiCallback    callback.Callback
 	uptime         time.Time
 
 	// mtx protects alertmanagerConfig, setAlertStatus and route.
@@ -86,16 +88,21 @@ func NewAPI(
 	gf groupsFn,
 	sf getAlertStatusFn,
 	silences *silence.Silences,
+	apiCallback callback.Callback,
 	peer cluster.ClusterPeer,
 	l log.Logger,
 	r prometheus.Registerer,
 ) (*API, error) {
+	if apiCallback == nil {
+		apiCallback = callback.NoopAPICallback{}
+	}
 	api := API{
 		alerts:         alerts,
 		getAlertStatus: sf,
 		alertGroups:    gf,
 		peer:           peer,
 		silences:       silences,
+		apiCallback:    apiCallback,
 		logger:         l,
 		m:              metrics.NewAlerts("v2", r),
 		uptime:         time.Now(),
@@ -302,7 +309,13 @@ func (api *API) getAlertsHandler(params alert_ops.GetAlertsParams) middleware.Re
 		return *res[i].Fingerprint < *res[j].Fingerprint
 	})
 
-	return alert_ops.NewGetAlertsOK().WithPayload(res)
+	callbackRes, err := api.apiCallback.V2GetAlertsCallback(res)
+	if err != nil {
+		level.Error(logger).Log("msg", "Failed to call api callback", "err", err)
+		return alert_ops.NewGetAlertsInternalServerError().WithPayload(err.Error())
+	}
+
+	return alert_ops.NewGetAlertsOK().WithPayload(callbackRes)
 }
 
 func (api *API) postAlertsHandler(params alert_ops.PostAlertsParams) middleware.Responder {
@@ -421,7 +434,13 @@ func (api *API) getAlertGroupsHandler(params alertgroup_ops.GetAlertGroupsParams
 		res = append(res, ag)
 	}
 
-	return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(res)
+	callbackRes, err := api.apiCallback.V2GetAlertGroupsCallback(res)
+	if err != nil {
+		level.Error(logger).Log("msg", "Failed to call api callback", "err", err)
+		return alertgroup_ops.NewGetAlertGroupsInternalServerError().WithPayload(err.Error())
+	}
+
+	return alertgroup_ops.NewGetAlertGroupsOK().WithPayload(callbackRes)
 }
 
 func (api *API) alertFilter(matchers []*labels.Matcher, silenced, inhibited, active bool) func(a *types.Alert, now time.Time) bool {

api/v2/api_test.go

@@ -15,14 +15,18 @@ package v2
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
 	"testing"
 	"time"
-
+	"github.com/prometheus/alertmanager/dispatch"
+	alert_ops "github.com/prometheus/alertmanager/api/v2/restapi/operations/alert"
+	alertgroup_ops "github.com/prometheus/alertmanager/api/v2/restapi/operations/alertgroup"
+	"github.com/prometheus/alertmanager/util/callback"
 	"github.com/go-openapi/runtime"
 	"github.com/go-openapi/strfmt"
 	"github.com/prometheus/common/model"
@@ -510,3 +514,222 @@ receivers:
 		require.Equal(t, tc.body, string(body))
 	}
 }
+
+func TestListAlertsHandler(t *testing.T) {
+	now := time.Now()
+	alerts := []*types.Alert{
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"alertname": "alert1"},
+				StartsAt: now.Add(-time.Minute),
+			},
+		},
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"alertname": "alert2"},
+				StartsAt: now.Add(-time.Minute),
+			},
+		},
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"alertname": "alert3"},
+				StartsAt: now.Add(-time.Minute),
+			},
+		},
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"alertname": "alert4"},
+				StartsAt: now.Add(-time.Minute),
+			},
+		},
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"alertname": "alert5"},
+				StartsAt: now.Add(-time.Minute),
+			},
+		},
+	}
+
+	for _, tc := range []struct {
+		name         string
+		expectedCode int
+		anames       []string
+		callback     callback.Callback
+	}{
+		{
+			"no call back",
+			200,
+			[]string{"alert3", "alert2", "alert1", "alert5", "alert4"},
+			callback.NoopAPICallback{},
+		},
+		{
+			"callback: only return 1 alert",
+			200,
+			[]string{"alert3"},
+			limitNumberOfAlertsReturnedCallback{limit: 1},
+		},
+		{
+			"callback: only return 3 alert",
+			200,
+			[]string{"alert3", "alert2", "alert1"},
+			limitNumberOfAlertsReturnedCallback{limit: 3},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			alertsProvider := newFakeAlerts(alerts)
+			api := API{
+				uptime:         time.Now(),
+				getAlertStatus: getAlertStatus,
+				logger:         log.NewNopLogger(),
+				apiCallback:    tc.callback,
+				alerts:         alertsProvider,
+				setAlertStatus: func(model.LabelSet) {},
+			}
+			api.route = dispatch.NewRoute(&config.Route{Receiver: "def-receiver"}, nil)
+			r, err := http.NewRequest("GET", "/api/v2/alerts", nil)
+			require.NoError(t, err)
+
+			w := httptest.NewRecorder()
+			p := runtime.TextProducer()
+			silence := false
+			inhibited := false
+			active := true
+			responder := api.getAlertsHandler(alert_ops.GetAlertsParams{
+				HTTPRequest: r,
+				Silenced:    &silence,
+				Inhibited:   &inhibited,
+				Active:      &active,
+			})
+			responder.WriteResponse(w, p)
+			body, _ := io.ReadAll(w.Result().Body)
+
+			require.Equal(t, tc.expectedCode, w.Code)
+			retAlerts := open_api_models.GettableAlerts{}
+			err = json.Unmarshal(body, &retAlerts)
+			if err != nil {
+				t.Fatalf("Unexpected error %v", err)
+			}
+			anames := []string{}
+			for _, a := range retAlerts {
+				name, ok := a.Labels["alertname"]
+				if ok {
+					anames = append(anames, string(name))
+				}
+			}
+			require.Equal(t, tc.anames, anames)
+		})
+	}
+}
+
+func TestGetAlertGroupsHandler(t *testing.T) {
+	var startAt time.Time
+	alerts := []*types.Alert{
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"state": "active", "alertname": "alert1"},
+				StartsAt: startAt,
+			},
+		},
+		{
+			Alert: model.Alert{
+				Labels:   model.LabelSet{"state": "unprocessed", "alertname": "alert2"},
+				StartsAt: startAt,
+			},
+		},
+	}
+	aginfos := dispatch.AlertGroups{
+		&dispatch.AlertGroup{
+			Labels: model.LabelSet{
+				"alertname": "TestingAlert",
+			},
+			Receiver: "testing",
+			Alerts:   alerts[:1],
+		},
+		&dispatch.AlertGroup{
+			Labels: model.LabelSet{
+				"alertname": "HighErrorRate",
+			},
+			Receiver: "prod",
+			Alerts:   alerts[:2],
+		},
+	}
+	for _, tc := range []struct {
+		name         string
+		numberOfAG   int
+		expectedCode int
+		callback     callback.Callback
+	}{
+		{
+			"no call back",
+			2,
+			200,
+			callback.NoopAPICallback{},
+		},
+		{
+			"callback: only return 1 alert group",
+			1,
+			200,
+			limitNumberOfAlertsReturnedCallback{limit: 1},
+		},
+		{
+			"callback: only return 2 alert group",
+			2,
+			200,
+			limitNumberOfAlertsReturnedCallback{limit: 2},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			api := API{
+				uptime: time.Now(),
+				alertGroups: func(func(*dispatch.Route) bool, func(*types.Alert, time.Time) bool) (dispatch.AlertGroups, map[model.Fingerprint][]string) {
+					return aginfos, nil
+				},
+				getAlertStatus: getAlertStatus,
+				logger:         log.NewNopLogger(),
+				apiCallback:    tc.callback,
+			}
+			r, err := http.NewRequest("GET", "/api/v2/alertgroups", nil)
+			require.NoError(t, err)
+
+			w := httptest.NewRecorder()
+			p := runtime.TextProducer()
+			silence := false
+			inhibited := false
+			active := true
+			responder := api.getAlertGroupsHandler(alertgroup_ops.GetAlertGroupsParams{
+				HTTPRequest: r,
+				Silenced:    &silence,
+				Inhibited:   &inhibited,
+				Active:      &active,
+			})
+			responder.WriteResponse(w, p)
+			body, _ := io.ReadAll(w.Result().Body)
+
+			require.Equal(t, tc.expectedCode, w.Code)
+			retAlertGroups := open_api_models.AlertGroups{}
+			err = json.Unmarshal(body, &retAlertGroups)
+			if err != nil {
+				t.Fatalf("Unexpected error %v", err)
+			}
+			require.Equal(t, tc.numberOfAG, len(retAlertGroups))
+		})
+	}
+}
+
+type limitNumberOfAlertsReturnedCallback struct {
+	limit int
+}
+
+func (n limitNumberOfAlertsReturnedCallback) V2GetAlertsCallback(alerts open_api_models.GettableAlerts) (open_api_models.GettableAlerts, error) {
+	return alerts[:n.limit], nil
+}
+
+func (n limitNumberOfAlertsReturnedCallback) V2GetAlertGroupsCallback(alertgroups open_api_models.AlertGroups) (open_api_models.AlertGroups, error) {
+	return alertgroups[:n.limit], nil
+}
+
+func getAlertStatus(model.Fingerprint) types.AlertStatus {
+	status := types.AlertStatus{SilencedBy: []string{}, InhibitedBy: []string{}}
+	status.State = types.AlertStateActive
+	return status
+}

api/v2/testing.go

@@ -18,6 +18,11 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/common/model"
+
+	"github.com/prometheus/alertmanager/provider"
+	"github.com/prometheus/alertmanager/types"
+
 	"github.com/go-openapi/strfmt"
 	"github.com/stretchr/testify/require"
 
@@ -68,3 +73,40 @@ func createLabelMatcher(t *testing.T, name, value string, matchType labels.Match
 	matcher, _ := labels.NewMatcher(matchType, name, value)
 	return matcher
 }
+
+// fakeAlerts is a struct implementing the provider.Alerts interface for tests.
+type fakeAlerts struct {
+	fps    map[model.Fingerprint]int
+	alerts []*types.Alert
+	err    error
+}
+
+func newFakeAlerts(alerts []*types.Alert) *fakeAlerts {
+	fps := make(map[model.Fingerprint]int)
+	for i, a := range alerts {
+		fps[a.Fingerprint()] = i
+	}
+	f := &fakeAlerts{
+		alerts: alerts,
+		fps:    fps,
+	}
+	return f
+}
+
+func (f *fakeAlerts) Subscribe() provider.AlertIterator           { return nil }
+func (f *fakeAlerts) Get(model.Fingerprint) (*types.Alert, error) { return nil, nil }
+func (f *fakeAlerts) Put(alerts ...*types.Alert) error {
+	return f.err
+}
+
+func (f *fakeAlerts) GetPending() provider.AlertIterator {
+	ch := make(chan *types.Alert)
+	done := make(chan struct{})
+	go func() {
+		defer close(ch)
+		for _, a := range f.alerts {
+			ch <- a
+		}
+	}()
+	return provider.NewAlertIterator(ch, done, f.err)
+}