From 72ecde50451b65ccb76b983c8b8968fcca7667fb Mon Sep 17 00:00:00 2001 From: Alexander Weaver Date: Fri, 2 Aug 2024 09:37:28 -0500 Subject: [PATCH] Alerting: Make orgID a direct arg of writer interface (#91422) make orgID a direct arg of writer interface --- pkg/services/ngalert/schedule/recording_rule.go | 2 +- pkg/services/ngalert/schedule/schedule.go | 2 +- pkg/services/ngalert/writer/fake.go | 6 +++--- pkg/services/ngalert/writer/noop.go | 2 +- pkg/services/ngalert/writer/prom.go | 10 ++-------- pkg/services/ngalert/writer/prom_test.go | 8 ++++---- 6 files changed, 12 insertions(+), 18 deletions(-) diff --git a/pkg/services/ngalert/schedule/recording_rule.go b/pkg/services/ngalert/schedule/recording_rule.go index 5d9391627ad..cb9878f40ea 100644 --- a/pkg/services/ngalert/schedule/recording_rule.go +++ b/pkg/services/ngalert/schedule/recording_rule.go @@ -257,7 +257,7 @@ func (r *recordingRule) tryEvaluation(ctx context.Context, ev *Evaluation, logge } writeStart := r.clock.Now() - err = r.writer.Write(ctx, ev.rule.Record.Metric, ev.scheduledAt, frames, ev.rule.Labels) + err = r.writer.Write(ctx, ev.rule.Record.Metric, ev.scheduledAt, frames, ev.rule.OrgID, ev.rule.Labels) writeDur := r.clock.Now().Sub(writeStart) if err != nil { diff --git a/pkg/services/ngalert/schedule/schedule.go b/pkg/services/ngalert/schedule/schedule.go index a8d3e038bce..4eeab054a57 100644 --- a/pkg/services/ngalert/schedule/schedule.go +++ b/pkg/services/ngalert/schedule/schedule.go @@ -49,7 +49,7 @@ type RulesStore interface { } type RecordingWriter interface { - Write(ctx context.Context, name string, t time.Time, frames data.Frames, extraLabels map[string]string) error + Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error } type schedule struct { diff --git a/pkg/services/ngalert/writer/fake.go b/pkg/services/ngalert/writer/fake.go index 44791cada02..bd688d65f7a 100644 --- a/pkg/services/ngalert/writer/fake.go +++ b/pkg/services/ngalert/writer/fake.go @@ -8,13 +8,13 @@ import ( ) type FakeWriter struct { - WriteFunc func(ctx context.Context, name string, t time.Time, frames data.Frames, extraLabels map[string]string) error + WriteFunc func(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error } -func (w FakeWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, extraLabels map[string]string) error { +func (w FakeWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error { if w.WriteFunc == nil { return nil } - return w.WriteFunc(ctx, name, t, frames, extraLabels) + return w.WriteFunc(ctx, name, t, frames, orgID, extraLabels) } diff --git a/pkg/services/ngalert/writer/noop.go b/pkg/services/ngalert/writer/noop.go index 3e434fdf76e..066f412fa1b 100644 --- a/pkg/services/ngalert/writer/noop.go +++ b/pkg/services/ngalert/writer/noop.go @@ -9,6 +9,6 @@ import ( type NoopWriter struct{} -func (w NoopWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, extraLabels map[string]string) error { +func (w NoopWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error { return nil } diff --git a/pkg/services/ngalert/writer/prom.go b/pkg/services/ngalert/writer/prom.go index ef860a599e7..bc4220b32bc 100644 --- a/pkg/services/ngalert/writer/prom.go +++ b/pkg/services/ngalert/writer/prom.go @@ -12,7 +12,6 @@ import ( "github.com/grafana/dataplane/sdata/numeric" "github.com/grafana/grafana/pkg/infra/log" "github.com/grafana/grafana/pkg/services/ngalert/metrics" - "github.com/grafana/grafana/pkg/services/ngalert/models" "github.com/grafana/grafana/pkg/setting" "github.com/m3db/prometheus_remote_client_golang/promremote" @@ -174,14 +173,9 @@ func createAuthOpts(username, password string) *httpclient.BasicAuthOptions { } // Write writes the given frames to the Prometheus remote write endpoint. -func (w PrometheusWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, extraLabels map[string]string) error { +func (w PrometheusWriter) Write(ctx context.Context, name string, t time.Time, frames data.Frames, orgID int64, extraLabels map[string]string) error { l := w.logger.FromContext(ctx) - ruleKey, found := models.RuleKeyFromContext(ctx) - if !found { - // sanity check, this should never happen - return fmt.Errorf("rule key not found in context") - } - lvs := []string{fmt.Sprint(ruleKey.OrgID), backendType} + lvs := []string{fmt.Sprint(orgID), backendType} points, err := PointsFromFrames(name, t, frames, extraLabels) if err != nil { diff --git a/pkg/services/ngalert/writer/prom_test.go b/pkg/services/ngalert/writer/prom_test.go index 020512b2e97..3aa49d85831 100644 --- a/pkg/services/ngalert/writer/prom_test.go +++ b/pkg/services/ngalert/writer/prom_test.go @@ -153,7 +153,7 @@ func TestPrometheusWriter_Write(t *testing.T) { ctx := ngmodels.WithRuleKey(context.Background(), ngmodels.GenerateRuleKey(1)) t.Run("error when frames are empty", func(t *testing.T) { - err := writer.Write(ctx, "test", now, emptyFrames, map[string]string{}) + err := writer.Write(ctx, "test", now, emptyFrames, 1, map[string]string{}) require.Error(t, err) }) @@ -163,7 +163,7 @@ func TestPrometheusWriter_Write(t *testing.T) { return promremote.WriteResult{}, clientErr } - err := writer.Write(ctx, "test", now, frames, map[string]string{}) + err := writer.Write(ctx, "test", now, frames, 1, map[string]string{}) require.Error(t, err) require.ErrorIs(t, err, clientErr) }) @@ -184,7 +184,7 @@ func TestPrometheusWriter_Write(t *testing.T) { return promremote.WriteResult{}, nil } - err := writer.Write(ctx, "test", now, frames, map[string]string{"extra": "label"}) + err := writer.Write(ctx, "test", now, frames, 1, map[string]string{"extra": "label"}) require.NoError(t, err) }) @@ -199,7 +199,7 @@ func TestPrometheusWriter_Write(t *testing.T) { return promremote.WriteResult{}, clientErr } - err := writer.Write(ctx, "test", now, frames, map[string]string{"extra": "label"}) + err := writer.Write(ctx, "test", now, frames, 1, map[string]string{"extra": "label"}) require.NoError(t, err) }) }