Alerting: Add traceability headers for alert queries (#57127)

* Define EvaluationContext

* Refactor ConditionEval to use new context struct

* Refactor QueriesAndExpressionsEval to use EvaluationContext

* Remove dead field from AlertExecCtx

* Refactor Validate to use EvaluationContext

* Get rid of privately used AlertExecCtx

* Move EvaluationContext to new file and add helper

* Add builder pattern and bind rule info to context

* Extract header logic and add rule UID header

* Fix missing call

pkg/services/ngalert/eval/context.go

@@ -0,0 +1,43 @@
+package eval
+
+import (
+	"context"
+	"time"
+
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/user"
+)
+
+// EvaluationContext represents the context in which a condition is evaluated.
+type EvaluationContext struct {
+	Ctx     context.Context
+	User    *user.SignedInUser
+	At      time.Time
+	RuleUID string
+}
+
+func Context(ctx context.Context, user *user.SignedInUser) EvaluationContext {
+	return EvaluationContext{
+		Ctx:  ctx,
+		User: user,
+		At:   time.Now(),
+	}
+}
+
+func (c EvaluationContext) When(t time.Time) EvaluationContext {
+	c.At = t
+	return c
+}
+
+func (c EvaluationContext) WithRule(r *models.AlertRule) EvaluationContext {
+	if r != nil {
+		c.RuleUID = r.UID
+	}
+	return c
+}
+
+func (c EvaluationContext) WithTimeout(timeout time.Duration) (EvaluationContext, context.CancelFunc) {
+	timeoutCtx, cancel := context.WithTimeout(c.Ctx, timeout)
+	c.Ctx = timeoutCtx
+	return c, cancel
+}

pkg/services/ngalert/eval/eval.go

@@ -3,7 +3,6 @@
 package eval
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"runtime/debug"
@@ -17,7 +16,6 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/ngalert/models"
-	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -27,11 +25,11 @@ import (
 //go:generate mockery --name Evaluator --structname FakeEvaluator --inpackage --filename evaluator_mock.go --with-expecter
 type Evaluator interface {
 	// ConditionEval executes conditions and evaluates the result.
-	ConditionEval(ctx context.Context, user *user.SignedInUser, condition models.Condition, now time.Time) Results
+	ConditionEval(ctx EvaluationContext, condition models.Condition) Results
 	// QueriesAndExpressionsEval executes queries and expressions and returns the result.
-	QueriesAndExpressionsEval(ctx context.Context, user *user.SignedInUser, data []models.AlertQuery, now time.Time) (*backend.QueryDataResponse, error)
+	QueriesAndExpressionsEval(ctx EvaluationContext, data []models.AlertQuery) (*backend.QueryDataResponse, error)
 	// Validate validates that the condition is correct. Returns nil if the condition is correct. Otherwise, error that describes the failure
-	Validate(ctx context.Context, user *user.SignedInUser, condition models.Condition) error
+	Validate(ctx EvaluationContext, condition models.Condition) error
 }
 
 type evaluatorImpl struct {
@@ -159,24 +157,31 @@ func (s State) String() string {
 	return [...]string{"Normal", "Alerting", "Pending", "NoData", "Error"}[s]
 }
 
-// AlertExecCtx is the context provided for executing an alert condition.
-type AlertExecCtx struct {
-	User               *user.SignedInUser
-	ExpressionsEnabled bool
-	Log                log.Logger
+func buildDatasourceHeaders(ctx EvaluationContext) map[string]string {
+	headers := map[string]string{
+		// Many data sources check this in query method as sometimes alerting needs special considerations.
+		// Several existing systems also compare against the value of this header. Altering this constitutes a breaking change.
+		//
+		// Note: The spelling of this headers is intentionally degenerate from the others for compatibility reasons.
+		// When sent over a network, the key of this header is canonicalized to "Fromalert".
+		// However, some datasources still compare against the string "FromAlert".
+		"FromAlert": "true",
 
-	Ctx context.Context
+		"X-Cache-Skip": "true",
+	}
+
+	if ctx.RuleUID != "" {
+		headers["X-Rule-Uid"] = ctx.RuleUID
+	}
+
+	return headers
 }
 
 // getExprRequest validates the condition, gets the datasource information and creates an expr.Request from it.
-func getExprRequest(ctx AlertExecCtx, data []models.AlertQuery, now time.Time, dsCacheService datasources.CacheService) (*expr.Request, error) {
+func getExprRequest(ctx EvaluationContext, data []models.AlertQuery, dsCacheService datasources.CacheService) (*expr.Request, error) {
 	req := &expr.Request{
-		OrgId: ctx.User.OrgID,
-		Headers: map[string]string{
-			// Some data sources check this in query method as sometimes alerting needs special considerations.
-			"FromAlert":    "true",
-			"X-Cache-Skip": "true",
-		},
+		OrgId:   ctx.User.OrgID,
+		Headers: buildDatasourceHeaders(ctx),
 	}
 
 	datasources := make(map[string]*datasources.DataSource, len(data))
@@ -211,8 +216,8 @@ func getExprRequest(ctx AlertExecCtx, data []models.AlertQuery, now time.Time, d
 
 		req.Queries = append(req.Queries, expr.Query{
 			TimeRange: expr.TimeRange{
-				From: q.RelativeTimeRange.ToTimeRange(now).From,
-				To:   q.RelativeTimeRange.ToTimeRange(now).To,
+				From: q.RelativeTimeRange.ToTimeRange(ctx.At).From,
+				To:   q.RelativeTimeRange.ToTimeRange(ctx.At).To,
 			},
 			DataSource:    ds,
 			JSON:          model,
@@ -311,10 +316,10 @@ func queryDataResponseToExecutionResults(c models.Condition, execResp *backend.Q
 	return result
 }
 
-func executeQueriesAndExpressions(ctx AlertExecCtx, data []models.AlertQuery, now time.Time, exprService *expr.Service, dsCacheService datasources.CacheService) (resp *backend.QueryDataResponse, err error) {
+func executeQueriesAndExpressions(ctx EvaluationContext, data []models.AlertQuery, exprService *expr.Service, dsCacheService datasources.CacheService, log log.Logger) (resp *backend.QueryDataResponse, err error) {
 	defer func() {
 		if e := recover(); e != nil {
-			ctx.Log.Error("alert rule panic", "error", e, "stack", string(debug.Stack()))
+			log.Error("alert rule panic", "error", e, "stack", string(debug.Stack()))
 			panicErr := fmt.Errorf("alert rule panic; please check the logs for the full stack")
 			if err != nil {
 				err = fmt.Errorf("queries and expressions execution failed: %w; %v", err, panicErr.Error())
@@ -324,7 +329,7 @@ func executeQueriesAndExpressions(ctx AlertExecCtx, data []models.AlertQuery, no
 		}
 	}()
 
-	queryDataReq, err := getExprRequest(ctx, data, now, dsCacheService)
+	queryDataReq, err := getExprRequest(ctx, data, dsCacheService)
 	if err != nil {
 		return nil, err
 	}
@@ -560,25 +565,23 @@ func (evalResults Results) AsDataFrame() data.Frame {
 }
 
 // ConditionEval executes conditions and evaluates the result.
-func (e *evaluatorImpl) ConditionEval(ctx context.Context, user *user.SignedInUser, condition models.Condition, now time.Time) Results {
-	execResp, err := e.QueriesAndExpressionsEval(ctx, user, condition.Data, now)
+func (e *evaluatorImpl) ConditionEval(ctx EvaluationContext, condition models.Condition) Results {
+	execResp, err := e.QueriesAndExpressionsEval(ctx, condition.Data)
 	var execResults ExecutionResults
 	if err != nil {
 		execResults = ExecutionResults{Error: err}
 	} else {
 		execResults = queryDataResponseToExecutionResults(condition, execResp)
 	}
-	return evaluateExecutionResult(execResults, now)
+	return evaluateExecutionResult(execResults, ctx.At)
 }
 
 // QueriesAndExpressionsEval executes queries and expressions and returns the result.
-func (e *evaluatorImpl) QueriesAndExpressionsEval(ctx context.Context, user *user.SignedInUser, data []models.AlertQuery, now time.Time) (*backend.QueryDataResponse, error) {
-	alertCtx, cancelFn := context.WithTimeout(ctx, e.cfg.UnifiedAlerting.EvaluationTimeout)
+func (e *evaluatorImpl) QueriesAndExpressionsEval(ctx EvaluationContext, data []models.AlertQuery) (*backend.QueryDataResponse, error) {
+	timeoutCtx, cancelFn := ctx.WithTimeout(e.cfg.UnifiedAlerting.EvaluationTimeout)
 	defer cancelFn()
 
-	alertExecCtx := AlertExecCtx{User: user, Ctx: alertCtx, ExpressionsEnabled: e.cfg.ExpressionsEnabled, Log: e.log}
-
-	execResult, err := executeQueriesAndExpressions(alertExecCtx, data, now, e.expressionService, e.dataSourceCache)
+	execResult, err := executeQueriesAndExpressions(timeoutCtx, data, e.expressionService, e.dataSourceCache, e.log)
 	if err != nil {
 		return nil, fmt.Errorf("failed to execute conditions: %w", err)
 	}
@@ -586,14 +589,7 @@ func (e *evaluatorImpl) QueriesAndExpressionsEval(ctx context.Context, user *use
 	return execResult, nil
 }
 
-func (e *evaluatorImpl) Validate(ctx context.Context, user *user.SignedInUser, condition models.Condition) error {
-	evalctx := AlertExecCtx{
-		User:               user,
-		ExpressionsEnabled: e.cfg.ExpressionsEnabled,
-		Log:                e.log,
-		Ctx:                ctx,
-	}
-
+func (e *evaluatorImpl) Validate(ctx EvaluationContext, condition models.Condition) error {
 	if len(condition.Data) == 0 {
 		return errors.New("expression list is empty. must be at least 1 expression")
 	}
@@ -601,7 +597,9 @@ func (e *evaluatorImpl) Validate(ctx context.Context, user *user.SignedInUser, c
 		return errors.New("condition must not be empty")
 	}
 
-	req, err := getExprRequest(evalctx, condition.Data, time.Now(), e.dataSourceCache)
+	ctx.At = time.Now()
+
+	req, err := getExprRequest(ctx, condition.Data, e.dataSourceCache)
 	if err != nil {
 		return err
 	}

pkg/services/ngalert/eval/eval_test.go

@@ -445,8 +445,9 @@ func TestValidate(t *testing.T) {
 			condition := testCase.condition(cacheService)
 
 			evaluator := NewEvaluator(&setting.Cfg{ExpressionsEnabled: true}, log.New("test"), cacheService, expr.ProvideService(&setting.Cfg{ExpressionsEnabled: true}, nil, nil))
+			evalCtx := Context(context.Background(), u)
 
-			err := evaluator.Validate(context.Background(), u, condition)
+			err := evaluator.Validate(evalCtx, condition)
 			if testCase.error {
 				require.Error(t, err)
 			} else {

pkg/services/ngalert/eval/evaluator_mock.go

@@ -1,19 +1,14 @@
-// Code generated by mockery v2.10.0. DO NOT EDIT.
+// Code generated by mockery v2.12.0. DO NOT EDIT.
 
 package eval
 
 import (
-	context "context"
-
 	backend "github.com/grafana/grafana-plugin-sdk-go/backend"
-
 	mock "github.com/stretchr/testify/mock"
 
 	models "github.com/grafana/grafana/pkg/services/ngalert/models"
 
-	time "time"
-
-	user "github.com/grafana/grafana/pkg/services/user"
+	testing "testing"
 )
 
 // FakeEvaluator is an autogenerated mock type for the Evaluator type
@@ -29,13 +24,13 @@ func (_m *FakeEvaluator) EXPECT() *FakeEvaluator_Expecter {
 	return &FakeEvaluator_Expecter{mock: &_m.Mock}
 }
 
-// ConditionEval provides a mock function with given fields: ctx, _a1, condition, now
-func (_m *FakeEvaluator) ConditionEval(ctx context.Context, _a1 *user.SignedInUser, condition models.Condition, now time.Time) Results {
-	ret := _m.Called(ctx, _a1, condition, now)
+// ConditionEval provides a mock function with given fields: ctx, condition
+func (_m *FakeEvaluator) ConditionEval(ctx EvaluationContext, condition models.Condition) Results {
+	ret := _m.Called(ctx, condition)
 
 	var r0 Results
-	if rf, ok := ret.Get(0).(func(context.Context, *user.SignedInUser, models.Condition, time.Time) Results); ok {
-		r0 = rf(ctx, _a1, condition, now)
+	if rf, ok := ret.Get(0).(func(EvaluationContext, models.Condition) Results); ok {
+		r0 = rf(ctx, condition)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(Results)
@@ -51,17 +46,15 @@ type FakeEvaluator_ConditionEval_Call struct {
 }
 
 // ConditionEval is a helper method to define mock.On call
-//   - ctx context.Context
-//   - _a1 *user.SignedInUser
-//   - condition models.Condition
-//   - now time.Time
-func (_e *FakeEvaluator_Expecter) ConditionEval(ctx interface{}, _a1 interface{}, condition interface{}, now interface{}) *FakeEvaluator_ConditionEval_Call {
-	return &FakeEvaluator_ConditionEval_Call{Call: _e.mock.On("ConditionEval", ctx, _a1, condition, now)}
+//  - ctx EvaluationContext
+//  - condition models.Condition
+func (_e *FakeEvaluator_Expecter) ConditionEval(ctx interface{}, condition interface{}) *FakeEvaluator_ConditionEval_Call {
+	return &FakeEvaluator_ConditionEval_Call{Call: _e.mock.On("ConditionEval", ctx, condition)}
 }
 
-func (_c *FakeEvaluator_ConditionEval_Call) Run(run func(ctx context.Context, _a1 *user.SignedInUser, condition models.Condition, now time.Time)) *FakeEvaluator_ConditionEval_Call {
+func (_c *FakeEvaluator_ConditionEval_Call) Run(run func(ctx EvaluationContext, condition models.Condition)) *FakeEvaluator_ConditionEval_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*user.SignedInUser), args[2].(models.Condition), args[3].(time.Time))
+		run(args[0].(EvaluationContext), args[1].(models.Condition))
 	})
 	return _c
 }
@@ -71,13 +64,13 @@ func (_c *FakeEvaluator_ConditionEval_Call) Return(_a0 Results) *FakeEvaluator_C
 	return _c
 }
 
-// QueriesAndExpressionsEval provides a mock function with given fields: ctx, _a1, data, now
-func (_m *FakeEvaluator) QueriesAndExpressionsEval(ctx context.Context, _a1 *user.SignedInUser, data []models.AlertQuery, now time.Time) (*backend.QueryDataResponse, error) {
-	ret := _m.Called(ctx, _a1, data, now)
+// QueriesAndExpressionsEval provides a mock function with given fields: ctx, data
+func (_m *FakeEvaluator) QueriesAndExpressionsEval(ctx EvaluationContext, data []models.AlertQuery) (*backend.QueryDataResponse, error) {
+	ret := _m.Called(ctx, data)
 
 	var r0 *backend.QueryDataResponse
-	if rf, ok := ret.Get(0).(func(context.Context, *user.SignedInUser, []models.AlertQuery, time.Time) *backend.QueryDataResponse); ok {
-		r0 = rf(ctx, _a1, data, now)
+	if rf, ok := ret.Get(0).(func(EvaluationContext, []models.AlertQuery) *backend.QueryDataResponse); ok {
+		r0 = rf(ctx, data)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*backend.QueryDataResponse)
@@ -85,8 +78,8 @@ func (_m *FakeEvaluator) QueriesAndExpressionsEval(ctx context.Context, _a1 *use
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(context.Context, *user.SignedInUser, []models.AlertQuery, time.Time) error); ok {
-		r1 = rf(ctx, _a1, data, now)
+	if rf, ok := ret.Get(1).(func(EvaluationContext, []models.AlertQuery) error); ok {
+		r1 = rf(ctx, data)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -100,17 +93,15 @@ type FakeEvaluator_QueriesAndExpressionsEval_Call struct {
 }
 
 // QueriesAndExpressionsEval is a helper method to define mock.On call
-//   - ctx context.Context
-//   - _a1 *user.SignedInUser
-//   - data []models.AlertQuery
-//   - now time.Time
-func (_e *FakeEvaluator_Expecter) QueriesAndExpressionsEval(ctx interface{}, _a1 interface{}, data interface{}, now interface{}) *FakeEvaluator_QueriesAndExpressionsEval_Call {
-	return &FakeEvaluator_QueriesAndExpressionsEval_Call{Call: _e.mock.On("QueriesAndExpressionsEval", ctx, _a1, data, now)}
+//  - ctx EvaluationContext
+//  - data []models.AlertQuery
+func (_e *FakeEvaluator_Expecter) QueriesAndExpressionsEval(ctx interface{}, data interface{}) *FakeEvaluator_QueriesAndExpressionsEval_Call {
+	return &FakeEvaluator_QueriesAndExpressionsEval_Call{Call: _e.mock.On("QueriesAndExpressionsEval", ctx, data)}
 }
 
-func (_c *FakeEvaluator_QueriesAndExpressionsEval_Call) Run(run func(ctx context.Context, _a1 *user.SignedInUser, data []models.AlertQuery, now time.Time)) *FakeEvaluator_QueriesAndExpressionsEval_Call {
+func (_c *FakeEvaluator_QueriesAndExpressionsEval_Call) Run(run func(ctx EvaluationContext, data []models.AlertQuery)) *FakeEvaluator_QueriesAndExpressionsEval_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*user.SignedInUser), args[2].([]models.AlertQuery), args[3].(time.Time))
+		run(args[0].(EvaluationContext), args[1].([]models.AlertQuery))
 	})
 	return _c
 }
@@ -120,13 +111,13 @@ func (_c *FakeEvaluator_QueriesAndExpressionsEval_Call) Return(_a0 *backend.Quer
 	return _c
 }
 
-// Validate provides a mock function with given fields: ctx, _a1, condition
-func (_m *FakeEvaluator) Validate(ctx context.Context, _a1 *user.SignedInUser, condition models.Condition) error {
-	ret := _m.Called(ctx, _a1, condition)
+// Validate provides a mock function with given fields: ctx, condition
+func (_m *FakeEvaluator) Validate(ctx EvaluationContext, condition models.Condition) error {
+	ret := _m.Called(ctx, condition)
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, *user.SignedInUser, models.Condition) error); ok {
-		r0 = rf(ctx, _a1, condition)
+	if rf, ok := ret.Get(0).(func(EvaluationContext, models.Condition) error); ok {
+		r0 = rf(ctx, condition)
 	} else {
 		r0 = ret.Error(0)
 	}
@@ -140,16 +131,15 @@ type FakeEvaluator_Validate_Call struct {
 }
 
 // Validate is a helper method to define mock.On call
-//   - ctx context.Context
-//   - _a1 *user.SignedInUser
-//   - condition models.Condition
-func (_e *FakeEvaluator_Expecter) Validate(ctx interface{}, _a1 interface{}, condition interface{}) *FakeEvaluator_Validate_Call {
-	return &FakeEvaluator_Validate_Call{Call: _e.mock.On("Validate", ctx, _a1, condition)}
+//  - ctx EvaluationContext
+//  - condition models.Condition
+func (_e *FakeEvaluator_Expecter) Validate(ctx interface{}, condition interface{}) *FakeEvaluator_Validate_Call {
+	return &FakeEvaluator_Validate_Call{Call: _e.mock.On("Validate", ctx, condition)}
 }
 
-func (_c *FakeEvaluator_Validate_Call) Run(run func(ctx context.Context, _a1 *user.SignedInUser, condition models.Condition)) *FakeEvaluator_Validate_Call {
+func (_c *FakeEvaluator_Validate_Call) Run(run func(ctx EvaluationContext, condition models.Condition)) *FakeEvaluator_Validate_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*user.SignedInUser), args[2].(models.Condition))
+		run(args[0].(EvaluationContext), args[1].(models.Condition))
 	})
 	return _c
 }
@@ -158,3 +148,13 @@ func (_c *FakeEvaluator_Validate_Call) Return(_a0 error) *FakeEvaluator_Validate
 	_c.Call.Return(_a0)
 	return _c
 }
+
+// NewFakeEvaluator creates a new instance of FakeEvaluator. It also registers the testing.TB interface on the mock and a cleanup function to assert the mocks expectations.
+func NewFakeEvaluator(t testing.TB) *FakeEvaluator {
+	mock := &FakeEvaluator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}