Autoinstrumentation improvements

backend/data_adapter.go

@@ -29,78 +29,54 @@ func (a *dataSDKAdapter) QueryData(ctx context.Context, req *pluginv2.QueryDataR
 		var innerErr error
 		resp, innerErr = a.queryDataHandler.QueryData(ctx, parsedReq)
 
-		if resp == nil || len(resp.Responses) == 0 {
-			return RequestStatusFromError(innerErr), innerErr
-		}
-
-		if isCancelledError(innerErr) {
-			return RequestStatusCancelled, nil
-		}
-
-		if isHTTPTimeoutError(innerErr) {
-			return RequestStatusError, nil
-		}
+		status := RequestStatusFromQueryDataResponse(resp, innerErr)
+		ctxLogger := Logger.FromContext(ctx)
 
 		// Set downstream status source in the context if there's at least one response with downstream status source,
 		// and if there's no plugin error
-		var hasPluginError bool
-		var hasDownstreamError bool
-		var hasCancelledError bool
-		var hasHTTPTimeoutError bool
-		for _, r := range resp.Responses {
+		var hasPluginError, hasDownstreamError bool
+		for refID, r := range resp.Responses {
 			if r.Error == nil {
 				continue
 			}
 
-			if isCancelledError(r.Error) {
-				hasCancelledError = true
+			// if error source not set and the error is a downstream error, set error source to downstream.
+			if !r.ErrorSource.IsValid() && IsDownstreamError(r.Error) {
+				r.ErrorSource = ErrorSourceDownstream
 			}
-			if isHTTPTimeoutError(r.Error) {
-				hasHTTPTimeoutError = true
+
+			if !r.Status.IsValid() {
+				r.Status = statusFromError(r.Error)
 			}
 
 			if r.ErrorSource == ErrorSourceDownstream {
 				hasDownstreamError = true
 			} else {
 				hasPluginError = true
 			}
-		}
-
-		if hasCancelledError {
-			if err := WithDownstreamErrorSource(ctx); err != nil {
-				return RequestStatusError, fmt.Errorf("failed to set downstream status source: %w", errors.Join(innerErr, err))
-			}
-			return RequestStatusCancelled, nil
-		}
 
-		if hasHTTPTimeoutError {
-			if err := WithDownstreamErrorSource(ctx); err != nil {
-				return RequestStatusError, fmt.Errorf("failed to set downstream status source: %w", errors.Join(innerErr, err))
+			logParams := []any{
+				"refID", refID,
+				"status", int(r.Status),
+				"error", r.Error,
+				"statusSource", string(r.ErrorSource),
 			}
-			return RequestStatusError, nil
+			ctxLogger.Error("Partial data response error", logParams...)
 		}
 
 		// A plugin error has higher priority than a downstream error,
 		// so set to downstream only if there's no plugin error
-		if hasDownstreamError && !hasPluginError {
-			if err := WithDownstreamErrorSource(ctx); err != nil {
-				return RequestStatusError, fmt.Errorf("failed to set downstream status source: %w", errors.Join(innerErr, err))
-			}
-			return RequestStatusError, nil
-		}
-
 		if hasPluginError {
 			if err := WithErrorSource(ctx, ErrorSourcePlugin); err != nil {
 				return RequestStatusError, fmt.Errorf("failed to set plugin status source: %w", errors.Join(innerErr, err))
 			}
-			return RequestStatusError, nil
-		}
-
-		if innerErr != nil {
-			return RequestStatusFromError(innerErr), innerErr
+		} else if hasDownstreamError {
+			if err := WithDownstreamErrorSource(ctx); err != nil {
+				return RequestStatusError, fmt.Errorf("failed to set downstream status source: %w", errors.Join(innerErr, err))
+			}
 		}
 
-		return RequestStatusOK, nil
+		return status, nil
 	})
 	if err != nil {
 		return nil, err

backend/data_adapter_test.go

@@ -180,6 +180,26 @@ func TestQueryData(t *testing.T) {
 				},
 				expErrorSource: ErrorSourcePlugin,
 			},
+			{
+				name: `single downstream error without error source should be "downstream" error source`,
+				queryDataResponse: &QueryDataResponse{
+					Responses: map[string]DataResponse{
+						"A": {Error: DownstreamErrorf("boom")},
+					},
+				},
+				expErrorSource: ErrorSourceDownstream,
+			},
+			{
+				name: `multiple downstream error without error source and single plugin error should be "plugin" error source`,
+				queryDataResponse: &QueryDataResponse{
+					Responses: map[string]DataResponse{
+						"A": {Error: DownstreamErrorf("boom")},
+						"B": {Error: someErr},
+						"C": {Error: DownstreamErrorf("boom")},
+					},
+				},
+				expErrorSource: ErrorSourcePlugin,
+			},
 		} {
 			t.Run(tc.name, func(t *testing.T) {
 				var actualCtx context.Context

backend/error_source.go

@@ -21,6 +21,10 @@ const (
 	DefaultErrorSource ErrorSource = ErrorSourcePlugin
 )
 
+func (es ErrorSource) IsValid() bool {
+	return es == ErrorSourceDownstream || es == ErrorSourcePlugin
+}
+
 // ErrorSourceFromStatus returns an [ErrorSource] based on provided HTTP status code.
 func ErrorSourceFromHTTPStatus(statusCode int) ErrorSource {
 	switch statusCode {
@@ -62,7 +66,7 @@ func IsDownstreamError(err error) bool {
 		return true
 	}
 
-	if isHTTPTimeoutError(err) {
+	if isHTTPTimeoutError(err) || isCancelledError(err) {
 		return true
 	}
 

backend/error_source_test.go

@@ -1,15 +1,26 @@
 package backend
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
 	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
+func TestErrorSource(t *testing.T) {
+	var es ErrorSource
+	require.False(t, es.IsValid())
+	require.True(t, ErrorSourceDownstream.IsValid())
+	require.True(t, ErrorSourcePlugin.IsValid())
+}
+
 func TestIsDownstreamError(t *testing.T) {
 	tcs := []struct {
 		name     string
@@ -66,6 +77,36 @@ func TestIsDownstreamError(t *testing.T) {
 			err:      fmt.Errorf("other error"),
 			expected: false,
 		},
+		{
+			name:     "context.Canceled",
+			err:      context.Canceled,
+			expected: true,
+		},
+		{
+			name:     "wrapped context.Canceled",
+			err:      fmt.Errorf("error: %w", context.Canceled),
+			expected: true,
+		},
+		{
+			name:     "joined context.Canceled",
+			err:      errors.Join(fmt.Errorf("oh no"), context.Canceled),
+			expected: true,
+		},
+		{
+			name:     "gRPC canceled error",
+			err:      status.Error(codes.Canceled, "canceled"),
+			expected: true,
+		},
+		{
+			name:     "wrapped gRPC canceled error",
+			err:      fmt.Errorf("error: %w", status.Error(codes.Canceled, "canceled")),
+			expected: true,
+		},
+		{
+			name:     "joined gRPC canceled error",
+			err:      errors.Join(fmt.Errorf("oh no"), status.Error(codes.Canceled, "canceled")),
+			expected: true,
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.name, func(t *testing.T) {

experimental/apis/data/v0alpha1/openapi_test.go

@@ -42,5 +42,5 @@ func TestOpenAPI(t *testing.T) {
 	// Ensure DataSourceRef exists and has three properties
 	def, ok = defs["github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1.DataSourceRef"]
 	require.True(t, ok)
-	require.Equal(t, []string{"type", "uid", "apiVersion"}, maps.Keys(def.Schema.Properties))
+	require.ElementsMatch(t, []string{"type", "uid", "apiVersion"}, maps.Keys(def.Schema.Properties))
 }