From 79300340af1c44fa12b483df2899aaa37b98cfa0 Mon Sep 17 00:00:00 2001
From: Danny Kopping <danny.kopping@grafana.com>
Date: Mon, 9 Jan 2023 10:14:37 +0200
Subject: [PATCH] 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       |  2 ++
 rules/alerting_test.go  | 32 ++++++++++++++++++++++++++++++++
 rules/origin.go         | 34 ++++++++++++++++++++++++++++++++++
 rules/recording.go      |  2 ++
 rules/recording_test.go | 22 ++++++++++++++++++++++
 5 files changed, 92 insertions(+)
 create mode 100644 rules/origin.go

diff --git a/rules/alerting.go b/rules/alerting.go
index d45666266..730e5a76f 100644
--- a/rules/alerting.go
+++ b/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
diff --git a/rules/alerting_test.go b/rules/alerting_test.go
index 4f5f5e683..bc41c0ff1 100644
--- a/rules/alerting_test.go
+++ b/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))
+}
diff --git a/rules/origin.go b/rules/origin.go
new file mode 100644
index 000000000..c972244a0
--- /dev/null
+++ b/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{}
+}
diff --git a/rules/recording.go b/rules/recording.go
index 8e7eefa1d..81ca47c91 100644
--- a/rules/recording.go
+++ b/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
diff --git a/rules/recording_test.go b/rules/recording_test.go
index 366dac52a..54c391215 100644
--- a/rules/recording_test.go
+++ b/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))
+}