Merge pull request #11826 from dannykopping/dannykopping/rule-eval

Pass rule details in evaluation context

rules/alerting.go

@@ -322,6 +322,8 @@ const resolvedRetention = 15 * time.Minute
 // Eval evaluates the rule expression and then creates pending alerts and fires
 // or removes previously pending alerts accordingly.
 func (r *AlertingRule) Eval(ctx context.Context, ts time.Time, query QueryFunc, externalURL *url.URL, limit int) (promql.Vector, error) {
+	ctx = NewOriginContext(ctx, NewRuleDetail(r))
+
 	res, err := query(ctx, r.vector.String(), ts)
 	if err != nil {
 		return nil, err

rules/alerting_test.go

@@ -895,3 +895,41 @@ func TestPendingAndKeepFiringFor(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 0, len(res))
 }
+
+// TestAlertingEvalWithOrigin checks that the alerting rule details are passed through the context.
+func TestAlertingEvalWithOrigin(t *testing.T) {
+	ctx := context.Background()
+	now := time.Now()
+
+	const (
+		name  = "my-recording-rule"
+		query = `count(metric{foo="bar"}) > 0`
+	)
+	var (
+		detail RuleDetail
+		lbs    = labels.FromStrings("test", "test")
+	)
+
+	expr, err := parser.ParseExpr(query)
+	require.NoError(t, err)
+
+	rule := NewAlertingRule(
+		name,
+		expr,
+		time.Second,
+		time.Minute,
+		lbs,
+		nil,
+		nil,
+		"",
+		true, log.NewNopLogger(),
+	)
+
+	_, err = rule.Eval(ctx, now, func(ctx context.Context, qs string, _ time.Time) (promql.Vector, error) {
+		detail = FromOriginContext(ctx)
+		return nil, nil
+	}, nil, 0)
+
+	require.NoError(t, err)
+	require.Equal(t, detail, NewRuleDetail(rule))
+}

rules/origin.go

@@ -0,0 +1,69 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rules
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/prometheus/prometheus/model/labels"
+)
+
+type ruleOrigin struct{}
+
+// RuleDetail contains information about the rule that is being evaluated.
+type RuleDetail struct {
+	Name   string
+	Query  string
+	Labels labels.Labels
+	Kind   string
+}
+
+const (
+	KindAlerting  = "alerting"
+	KindRecording = "recording"
+)
+
+// NewRuleDetail creates a RuleDetail from a given Rule.
+func NewRuleDetail(r Rule) RuleDetail {
+	var kind string
+	switch r.(type) {
+	case *AlertingRule:
+		kind = KindAlerting
+	case *RecordingRule:
+		kind = KindRecording
+	default:
+		panic(fmt.Sprintf(`unknown rule type "%T"`, r))
+	}
+
+	return RuleDetail{
+		Name:   r.Name(),
+		Query:  r.Query().String(),
+		Labels: r.Labels(),
+		Kind:   kind,
+	}
+}
+
+// NewOriginContext returns a new context with data about the origin attached.
+func NewOriginContext(ctx context.Context, rule RuleDetail) context.Context {
+	return context.WithValue(ctx, ruleOrigin{}, rule)
+}
+
+// FromOriginContext returns the RuleDetail origin data from the context.
+func FromOriginContext(ctx context.Context) RuleDetail {
+	if rule, ok := ctx.Value(ruleOrigin{}).(RuleDetail); ok {
+		return rule
+	}
+	return RuleDetail{}
+}

rules/origin_test.go

@@ -0,0 +1,51 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rules
+
+import (
+	"context"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
+)
+
+type unknownRule struct{}
+
+func (u unknownRule) Name() string          { return "" }
+func (u unknownRule) Labels() labels.Labels { return nil }
+func (u unknownRule) Eval(ctx context.Context, time time.Time, queryFunc QueryFunc, url *url.URL, i int) (promql.Vector, error) {
+	return nil, nil
+}
+func (u unknownRule) String() string                               { return "" }
+func (u unknownRule) Query() parser.Expr                           { return nil }
+func (u unknownRule) SetLastError(err error)                       {}
+func (u unknownRule) LastError() error                             { return nil }
+func (u unknownRule) SetHealth(health RuleHealth)                  {}
+func (u unknownRule) Health() RuleHealth                           { return "" }
+func (u unknownRule) SetEvaluationDuration(duration time.Duration) {}
+func (u unknownRule) GetEvaluationDuration() time.Duration         { return 0 }
+func (u unknownRule) SetEvaluationTimestamp(time time.Time)        {}
+func (u unknownRule) GetEvaluationTimestamp() time.Time            { return time.Time{} }
+
+func TestNewRuleDetailPanics(t *testing.T) {
+	require.PanicsWithValue(t, `unknown rule type "rules.unknownRule"`, func() {
+		NewRuleDetail(unknownRule{})
+	})
+}

rules/recording.go

@@ -73,6 +73,8 @@ func (rule *RecordingRule) Labels() labels.Labels {
 
 // Eval evaluates the rule and then overrides the metric names and labels accordingly.
 func (rule *RecordingRule) Eval(ctx context.Context, ts time.Time, query QueryFunc, _ *url.URL, limit int) (promql.Vector, error) {
+	ctx = NewOriginContext(ctx, NewRuleDetail(rule))
+
 	vector, err := query(ctx, rule.vector.String(), ts)
 	if err != nil {
 		return nil, err

rules/recording_test.go

@@ -156,3 +156,31 @@ func TestRecordingRuleLimit(t *testing.T) {
 		}
 	}
 }
+
+// TestRecordingEvalWithOrigin checks that the recording rule details are passed through the context.
+func TestRecordingEvalWithOrigin(t *testing.T) {
+	ctx := context.Background()
+	now := time.Now()
+
+	const (
+		name  = "my-recording-rule"
+		query = `count(metric{foo="bar"})`
+	)
+
+	var (
+		detail RuleDetail
+		lbs    = labels.FromStrings("foo", "bar")
+	)
+
+	expr, err := parser.ParseExpr(query)
+	require.NoError(t, err)
+
+	rule := NewRecordingRule(name, expr, lbs)
+	_, err = rule.Eval(ctx, now, func(ctx context.Context, qs string, _ time.Time) (promql.Vector, error) {
+		detail = FromOriginContext(ctx)
+		return nil, nil
+	}, nil, 0)
+
+	require.NoError(t, err)
+	require.Equal(t, detail, NewRuleDetail(rule))
+}