From 2eab7cb63b7127f5acf02d788800d018fbbd3847 Mon Sep 17 00:00:00 2001
From: Adam Simpson <adam@adamsimpson.net>
Date: Tue, 22 Apr 2025 15:20:45 -0400
Subject: [PATCH] ds-querier: improve instant vector support (#103954)

---
 pkg/registry/apis/query/query.go      | 14 +++----
 pkg/registry/apis/query/query_test.go | 56 +++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/pkg/registry/apis/query/query.go b/pkg/registry/apis/query/query.go
index da73b222fab..4eaa024c4f6 100644
--- a/pkg/registry/apis/query/query.go
+++ b/pkg/registry/apis/query/query.go
@@ -191,9 +191,9 @@ func (b *QueryAPIBuilder) execute(ctx context.Context, req parsedRequestInfo) (q
 	case 1:
 		b.log.Debug("executing single query")
 		qdr, err = b.handleQuerySingleDatasource(ctx, req.Requests[0])
-		if err == nil && alertQueryWithoutExpression(req) {
-			b.log.Debug("handling alert query without expression")
-			qdr, err = b.convertQueryWithoutExpression(ctx, req.Requests[0], qdr)
+		if err == nil && isSingleAlertQuery(req) {
+			b.log.Debug("handling alert query with single query")
+			qdr, err = b.convertQueryFromAlerting(ctx, req.Requests[0], qdr)
 		}
 	default:
 		b.log.Debug("executing concurrent queries")
@@ -416,7 +416,7 @@ func (b *QueryAPIBuilder) handleExpressions(ctx context.Context, req parsedReque
 	return qdr, nil
 }
 
-func (b *QueryAPIBuilder) convertQueryWithoutExpression(ctx context.Context, req datasourceRequest,
+func (b *QueryAPIBuilder) convertQueryFromAlerting(ctx context.Context, req datasourceRequest,
 	qdr *backend.QueryDataResponse) (*backend.QueryDataResponse, error) {
 	if len(req.Request.Queries) == 0 {
 		return nil, errors.New("no queries to convert")
@@ -475,14 +475,14 @@ func (r responderWrapper) Error(err error) {
 	r.wrapped.Error(err)
 }
 
-// Checks if the request only contains a single query and not expression.
-func alertQueryWithoutExpression(req parsedRequestInfo) bool {
+// Checks if the request only contains a single query and is from Alerting
+func isSingleAlertQuery(req parsedRequestInfo) bool {
 	if len(req.Requests) != 1 {
 		return false
 	}
 	headers := req.Requests[0].Headers
 	_, exist := headers[models.FromAlertHeaderName]
-	if exist && len(req.Requests[0].Request.Queries) == 1 && len(req.Expressions) == 0 {
+	if exist && len(req.Requests[0].Request.Queries) == 1 {
 		return true
 	}
 	return false
diff --git a/pkg/registry/apis/query/query_test.go b/pkg/registry/apis/query/query_test.go
index 5147b51bfb9..fbe1a2766e7 100644
--- a/pkg/registry/apis/query/query_test.go
+++ b/pkg/registry/apis/query/query_test.go
@@ -7,8 +7,10 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	frameData "github.com/grafana/grafana-plugin-sdk-go/data"
 	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
 	"github.com/grafana/grafana/pkg/expr"
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -81,6 +83,60 @@ func TestQueryRestConnectHandler(t *testing.T) {
 	}, *b.client.(mockClient).lastCalledWithHeaders)
 }
 
+func TestInstantQueryFromAlerting(t *testing.T) {
+	builder := &QueryAPIBuilder{
+		converter: &expr.ResultConverter{
+			Features: featuremgmt.WithFeatures(),
+			Tracer:   tracing.InitializeTracerForTest(),
+		},
+	}
+
+	dq := data.DataQuery{}
+	dq.RefID = "A"
+
+	dr := datasourceRequest{
+		Headers: map[string]string{
+			models.FromAlertHeaderName: "true",
+		},
+		Request: &data.QueryDataRequest{
+			Queries: []data.DataQuery{
+				dq,
+			},
+		},
+	}
+
+	fakeFrame := frameData.NewFrame(
+		"A",
+		frameData.NewField("Time", nil, []time.Time{time.Now()}),
+		frameData.NewField("Value", nil, []int64{42}),
+	)
+	fakeFrame.Meta = &frameData.FrameMeta{TypeVersion: frameData.FrameTypeVersion{0, 1}, Type: "numeric-multi"}
+
+	inputQDR := &backend.QueryDataResponse{
+		Responses: map[string]backend.DataResponse{
+			"A": {
+				Frames: frameData.Frames{
+					fakeFrame,
+				},
+			},
+		},
+	}
+
+	request := parsedRequestInfo{
+		Requests: []datasourceRequest{
+			dr,
+		},
+	}
+
+	result, err := builder.convertQueryFromAlerting(context.Background(), dr, inputQDR)
+	require.NoError(t, err)
+
+	require.True(t, isSingleAlertQuery(request), "Expected a valid alert query with a single query to return true")
+	require.NotNil(t, result)
+	require.Equal(t, 1, len(result.Responses["A"].Frames[0].Fields), "Expected a single field not Time and Value")
+	require.Equal(t, "Value", result.Responses["A"].Frames[0].Fields[0].Name, "Expected the single field to be Value")
+}
+
 type mockResponder struct {
 }