Add Rule Name and Type to Query Tags

grafana · Nov 21, 2024 · aab3185 · aab3185
1 parent 469b54e
commit aab3185
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 1 deletion.
diff --git a/pkg/ruler/compat.go b/pkg/ruler/compat.go
@@ -72,6 +72,12 @@ func queryFunc(evaluator Evaluator, checker readyChecker, userID string, logger
 			return nil, errNotReady
 		}
 
+		// Extract rule details
+		ruleName := detail.Name
+		ruleType := detail.Kind
+
+		// Add rule details to context
+		ctx = AddRuleDetailsToContext(ctx, ruleName, ruleType)
 		res, err := evaluator.Eval(ctx, qs, t)
 
 		if err != nil {
@@ -357,3 +363,25 @@ type noopRuleDependencyController struct{}
 func (*noopRuleDependencyController) AnalyseRules([]rules.Rule) {
 	// Do nothing
 }
+
+// Define context keys to avoid collisions
+type contextKey string
+
+const (
+	ruleNameKey contextKey = "rule_name"
+	ruleTypeKey contextKey = "rule_type"
+)
+
+// AddRuleDetailsToContext adds rule details to the context
+func AddRuleDetailsToContext(ctx context.Context, ruleName string, ruleType string) context.Context {
+	ctx = context.WithValue(ctx, ruleNameKey, ruleName)
+	ctx = context.WithValue(ctx, ruleTypeKey, ruleType)
+	return ctx
+}
+
+// GetRuleDetailsFromContext retrieves rule details from the context
+func GetRuleDetailsFromContext(ctx context.Context) (string, string) {
+	ruleName, _ := ctx.Value(ruleNameKey).(string)
+	ruleType, _ := ctx.Value(ruleTypeKey).(string)
+	return ruleName, ruleType
+}
diff --git a/pkg/ruler/compat_test.go b/pkg/ruler/compat_test.go
@@ -130,3 +130,19 @@ type fakeChecker struct{}
 func (f fakeChecker) isReady(_ string) bool {
 	return true
 }
+
+func TestAddAndGetRuleDetailsFromContext(t *testing.T) {
+	ctx := context.Background()
+	ruleName := "test_rule"
+	ruleType := "test_type"
+
+	// Add rule details to context
+	ctx = AddRuleDetailsToContext(ctx, ruleName, ruleType)
+
+	// Retrieve rule details from context
+	retrievedRuleName, retrievedRuleType := GetRuleDetailsFromContext(ctx)
+
+	// Assert that the retrieved values match the expected values
+	assert.Equal(t, ruleName, retrievedRuleName, "Expected rule name to match")
+	assert.Equal(t, ruleType, retrievedRuleType, "Expected rule type to match")
+}
diff --git a/pkg/ruler/evaluator_remote.go b/pkg/ruler/evaluator_remote.go
@@ -223,6 +223,11 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
 	body := []byte(args.Encode())
 	hash := util.HashedQuery(query)
 
+	// Retrieve rule details from context
+	ruleName, ruleType := GetRuleDetailsFromContext(ctx)
+
+	// Construct the X-Query-Tags header value
+	queryTags := fmt.Sprintf("source=ruler,rule_name=%s,rule_type=%s", ruleName, ruleType)
 	req := httpgrpc.HTTPRequest{
 		Method: http.MethodPost,
 		Url:    queryEndpointPath,
@@ -231,7 +236,7 @@ func (r *RemoteEvaluator) query(ctx context.Context, orgID, query string, ts tim
 			{Key: textproto.CanonicalMIMEHeaderKey("User-Agent"), Values: []string{userAgent}},
 			{Key: textproto.CanonicalMIMEHeaderKey("Content-Type"), Values: []string{mimeTypeFormPost}},
 			{Key: textproto.CanonicalMIMEHeaderKey("Content-Length"), Values: []string{strconv.Itoa(len(body))}},
-			{Key: textproto.CanonicalMIMEHeaderKey(string(httpreq.QueryTagsHTTPHeader)), Values: []string{"source=ruler"}},
+			{Key: textproto.CanonicalMIMEHeaderKey(string(httpreq.QueryTagsHTTPHeader)), Values: []string{queryTags}},
 			{Key: textproto.CanonicalMIMEHeaderKey(user.OrgIDHeaderName), Values: []string{orgID}},
 		},
 	}

diff --git a/pkg/ruler/evaluator_remote_test.go b/pkg/ruler/evaluator_remote_test.go
@@ -61,6 +61,7 @@ func TestRemoteEvalQueryTimeout(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.Error(t, err)
@@ -98,6 +99,7 @@ func TestRemoteEvalMaxResponseSize(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.Error(t, err)
@@ -146,6 +148,7 @@ func TestRemoteEvalScalar(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	res, err := ev.Eval(ctx, "19", now)
 	require.NoError(t, err)
@@ -189,6 +192,7 @@ func TestRemoteEvalEmptyScalarResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	res, err := ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.NoError(t, err)
@@ -247,6 +251,7 @@ func TestRemoteEvalVectorResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	res, err := ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", now)
 	require.NoError(t, err)
@@ -294,6 +299,7 @@ func TestRemoteEvalEmptyVectorResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.NoError(t, err)
@@ -317,6 +323,7 @@ func TestRemoteEvalErrorResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.ErrorContains(t, err, "rule evaluation failed")
@@ -343,6 +350,7 @@ func TestRemoteEvalNon2xxResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.ErrorContains(t, err, fmt.Sprintf("unsuccessful/unexpected response - status code %d", httpErr))
@@ -367,6 +375,7 @@ func TestRemoteEvalNonJSONResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.ErrorContains(t, err, "unexpected body encoding, not valid JSON")
@@ -406,6 +415,7 @@ func TestRemoteEvalUnsupportedResultResponse(t *testing.T) {
 
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "test")
+	ctx = AddRuleDetailsToContext(ctx, "test_rule_name", "test_rule_type")
 
 	_, err = ev.Eval(ctx, "sum(rate({foo=\"bar\"}[5m]))", time.Now())
 	require.ErrorContains(t, err, fmt.Sprintf("unsupported result type: %q", loghttp.ResultTypeStream))