Adding recording/alerting rule origin context

This will allow correlation of executed rule queries with their associated rule names and type

Signed-off-by: Danny Kopping <danny.kopping@grafana.com>

rules/alerting.go

@@ -311,6 +311,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.Name(), r.Query().String(), KindAlerting))
+
 	res, err := query(ctx, r.vector.String(), ts)
 	if err != nil {
 		return nil, err

rules/alerting_test.go

@@ -712,3 +712,35 @@ func TestSendAlertsDontAffectActiveAlerts(t *testing.T) {
 	// But the labels with the AlertingRule should not be changed.
 	require.Equal(t, labels.FromStrings("a1", "1"), rule.active[h].Labels)
 }
+
+func TestAlertingEvalWithOrigin(t *testing.T) {
+	ctx := context.Background()
+	now := time.Now()
+
+	const name = "my-recording-rule"
+	const query = `count(metric{foo="bar"}) > 0`
+
+	var detail RuleDetail
+
+	expr, err := parser.ParseExpr(query)
+	require.NoError(t, err)
+
+	rule := NewAlertingRule(
+		name,
+		expr,
+		time.Minute,
+		labels.FromStrings("test", "test"),
+		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(name, query, KindAlerting))
+}

rules/origin.go

@@ -0,0 +1,34 @@
+package rules
+
+import "context"
+
+type ruleOrigin struct{}
+
+type RuleDetail struct {
+	Name  string
+	Kind  string
+	Query string
+}
+
+const KindAlerting = "alerting"
+const KindRecording = "recording"
+
+func NewRuleDetail(name, query, kind string) RuleDetail {
+	return RuleDetail{
+		Name:  name,
+		Query: query,
+		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)
+}
+
+func FromOriginContext(ctx context.Context) RuleDetail {
+	if rule, ok := ctx.Value(ruleOrigin{}).(RuleDetail); ok {
+		return rule
+	}
+	return RuleDetail{}
+}

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.Name(), rule.Query().String(), KindRecording))
+
 	vector, err := query(ctx, rule.vector.String(), ts)
 	if err != nil {
 		return nil, err

rules/recording_test.go

@@ -156,3 +156,25 @@ func TestRecordingRuleLimit(t *testing.T) {
 		}
 	}
 }
+
+func TestRecordingEvalWithOrigin(t *testing.T) {
+	ctx := context.Background()
+	now := time.Now()
+
+	const name = "my-recording-rule"
+	const query = `count(metric{foo="bar"})`
+
+	var detail RuleDetail
+
+	expr, err := parser.ParseExpr(query)
+	require.NoError(t, err)
+
+	rule := NewRecordingRule(name, expr, []labels.Label{})
+	_, 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(name, query, KindRecording))
+}