query-frontend: correctly replace @ modifier for split queries

The existing code would only evaluate the at modifier for vector selectors, but subquery expressions also can have it.

Signed-off-by: Dimitar Dimitrov <dimitar.dimitrov@grafana.com>

pkg/frontend/querymiddleware/astmapper/instant_splitting_test.go

@@ -605,6 +605,10 @@ func TestInstantSplitterSkippedQueryReason(t *testing.T) {
 			query:         `max_over_time(absent_over_time(deriv(rate(metric_counter[1m])[5m:1m])[2m:])[10m:])`,
 			skippedReason: SkippedReasonSubquery,
 		},
+		{
+			query:         `sum by(group_1) (sum_over_time(metric_counter[7d:] @ start()))`,
+			skippedReason: SkippedReasonSubquery,
+		},
 	} {
 		tt := tt
 

pkg/frontend/querymiddleware/split_and_cache.go

@@ -667,7 +667,16 @@ func evaluateAtModifierFunction(query string, start, end int64) (string, error)
 		return "", apierror.New(apierror.TypeBadData, decorateWithParamName(err, "query").Error())
 	}
 	parser.Inspect(expr, func(n parser.Node, _ []parser.Node) error {
-		if selector, ok := n.(*parser.VectorSelector); ok {
+		switch selector := n.(type) {
+		case *parser.VectorSelector:
+			switch selector.StartOrEnd {
+			case parser.START:
+				selector.Timestamp = &start
+			case parser.END:
+				selector.Timestamp = &end
+			}
+			selector.StartOrEnd = 0
+		case *parser.SubqueryExpr:
 			switch selector.StartOrEnd {
 			case parser.START:
 				selector.Timestamp = &start

pkg/frontend/querymiddleware/split_and_cache_test.go

@@ -1620,6 +1620,10 @@ func TestSplitQueryByInterval(t *testing.T) {
 	queryFooAtStartExpr, _ := parser.ParseExpr(queryFooAtStart)
 	queryFooAtZero := "foo @ 0.000"
 	queryFooAtZeroExpr, _ := parser.ParseExpr(queryFooAtZero)
+	queryFooSubqueryAtStart := "sum_over_time(foo[1d:] @ start())"
+	queryFooSubqueryAtStartExpr, _ := parser.ParseExpr(queryFooSubqueryAtStart)
+	queryFooSubqueryAtZero := "sum_over_time(foo[1d:] @ 0.000)"
+	queryFooSubqueryAtZeroExpr, _ := parser.ParseExpr(queryFooSubqueryAtZero)
 
 	for i, tc := range []struct {
 		input    MetricsQueryRequest
@@ -1662,6 +1666,14 @@ func TestSplitQueryByInterval(t *testing.T) {
 			},
 			interval: day,
 		},
+		{
+			input: &PrometheusRangeQueryRequest{minT: -(24 * 3600 * seconds), start: 0, end: 2 * 24 * 3600 * seconds, step: 15 * seconds, queryExpr: queryFooSubqueryAtStartExpr},
+			expected: []MetricsQueryRequest{
+				&PrometheusRangeQueryRequest{minT: -(24 * 3600 * seconds), start: 0, end: (24 * 3600 * seconds) - (15 * seconds), step: 15 * seconds, queryExpr: queryFooSubqueryAtZeroExpr},
+				&PrometheusRangeQueryRequest{minT: -(24 * 3600 * seconds), start: 24 * 3600 * seconds, end: 2 * 24 * 3600 * seconds, step: 15 * seconds, queryExpr: queryFooSubqueryAtZeroExpr},
+			},
+			interval: day,
+		},
 		{
 			input: &PrometheusRangeQueryRequest{start: 0, end: 2 * 3 * 3600 * seconds, step: 15 * seconds, queryExpr: queryFooExpr},
 			expected: []MetricsQueryRequest{
@@ -1797,6 +1809,8 @@ func Test_evaluateAtModifier(t *testing.T) {
 		{"topk(5, rate(http_requests_total[1h] @ start()))", "topk(5, rate(http_requests_total[1h] @ 1546300.800))", nil},
 		{"topk(5, rate(http_requests_total[1h] @ 0))", "topk(5, rate(http_requests_total[1h] @ 0.000))", nil},
 		{"http_requests_total[1h] @ 10.001", "http_requests_total[1h] @ 10.001", nil},
+		{"sum_over_time(http_requests_total[1h:] @ start())", "sum_over_time(http_requests_total[1h:] @ 1546300.800)", nil},
+		{"sum_over_time((http_requests_total @ end())[1h:] @ start())", "sum_over_time((http_requests_total @ 1646300.800)[1h:] @ 1546300.800)", nil},
 		{
 			`min_over_time(
 				sum by(cluster) (