From 333a21f19de5ed7038fab81f428eba39a1709bca Mon Sep 17 00:00:00 2001
From: Sarah Zinger <sarah.zinger@grafana.com>
Date: Fri, 22 Aug 2025 14:59:07 -0400
Subject: [PATCH] DS-Querier: Fix data source not found (#110064)

---
 pkg/registry/apis/query/query.go | 39 ++++++++++++++++++++++++--------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/pkg/registry/apis/query/query.go b/pkg/registry/apis/query/query.go
index 25bf6439908..4539662c7eb 100644
--- a/pkg/registry/apis/query/query.go
+++ b/pkg/registry/apis/query/query.go
@@ -167,7 +167,7 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
 		raw := &query.QueryDataRequest{}
 		err := web.Bind(httpreq, raw)
 		if err != nil {
-			b.log.Error("Hit unexpected error when reading query", "err", err)
+			connectLogger.Error("Hit unexpected error when reading query", "err", err)
 			err = errorsK8s.NewBadRequest("error reading query")
 			// TODO: can we wrap the error so details are not lost?!
 			// errutil.BadRequest(
@@ -181,8 +181,8 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
 		qdr, err := handleQuery(ctx, *raw, *b, httpreq, *responder, connectLogger)
 
 		if err != nil {
-			b.log.Error("execute error", "http code", query.GetResponseCode(qdr), "err", err)
-			logEmptyRefids(raw.Queries, b.log)
+			connectLogger.Error("execute error", "http code", query.GetResponseCode(qdr), "err", err)
+			logEmptyRefids(raw.Queries, connectLogger)
 			if qdr != nil { // if we have a response, we assume the err is set in the response
 				responder.Object(query.GetResponseCode(qdr), &query.QueryDataResponse{
 					QueryDataResponse: *qdr,
@@ -190,19 +190,40 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
 				return
 			} else {
 				var errorDataResponse backend.DataResponse
-				if errors.Is(err, service.ErrInvalidDatasourceID) || errors.Is(err, service.ErrNoQueriesFound) || errors.Is(err, service.ErrMissingDataSourceInfo) || errors.Is(err, service.ErrQueryParamMismatch) || errors.Is(err, service.ErrDuplicateRefId) {
+
+				badRequestErrors := []error{
+					service.ErrInvalidDatasourceID,
+					service.ErrNoQueriesFound,
+					service.ErrMissingDataSourceInfo,
+					service.ErrQueryParamMismatch,
+					service.ErrDuplicateRefId,
+					datasources.ErrDataSourceNotFound,
+				}
+				isTypedBadRequestError := false
+				for _, badRequestError := range badRequestErrors {
+					if errors.Is(err, badRequestError) {
+						isTypedBadRequestError = true
+					}
+				}
+				if isTypedBadRequestError {
 					errorDataResponse = backend.ErrDataResponseWithSource(backend.StatusBadRequest, backend.ErrorSourceDownstream, err.Error())
 				} else if strings.Contains(err.Error(), "expression request error") {
-					b.log.Error("Error calling TransformData in an expression", "err", err)
+					connectLogger.Error("Error calling TransformData in an expression", "err", err)
 					errorDataResponse = backend.ErrDataResponseWithSource(backend.StatusBadRequest, backend.ErrorSourceDownstream, err.Error())
 				} else {
-					b.log.Error("unknown error, treated as a 500", "err", err)
+					connectLogger.Error("unknown error, treated as a 500", "err", err)
 					responder.Error(err)
 					return
 				}
+				// TODO ensure errors also return the refId wherever possible
+				errorRefId := raw.Queries[0].RefID
+				if errorRefId == "" {
+					errorRefId = "A"
+				}
+
 				qdr = &backend.QueryDataResponse{
 					Responses: map[string]backend.DataResponse{
-						"A": errorDataResponse,
+						errorRefId: errorDataResponse,
 					},
 				}
 				responder.Object(query.GetResponseCode(qdr), &query.QueryDataResponse{
@@ -223,12 +244,12 @@ func handleQuery(ctx context.Context, raw query.QueryDataRequest, b QueryAPIBuil
 	for _, query := range raw.Queries {
 		jsonBytes, err := json.Marshal(query)
 		if err != nil {
-			b.log.Error("error marshalling", err)
+			connectLogger.Error("error marshalling", err)
 		}
 
 		sjQuery, _ := simplejson.NewJson(jsonBytes)
 		if err != nil {
-			b.log.Error("error unmarshalling", err)
+			connectLogger.Error("error unmarshalling", err)
 		}
 
 		jsonQueries = append(jsonQueries, sjQuery)