Error source handler middleware

backend/adapter_utils.go

@@ -75,7 +75,7 @@ func metricWrapper(next handlerWrapperFunc) handlerWrapperFunc {
 		endpoint := EndpointFromContext(ctx)
 		status, err := next(ctx)
 
-		pluginRequestCounter.WithLabelValues(endpoint.String(), status.String(), string(errorSourceFromContext(ctx))).Inc()
+		pluginRequestCounter.WithLabelValues(endpoint.String(), status.String(), string(ErrorSourceFromContext(ctx))).Inc()
 
 		return status, err
 	}
@@ -106,7 +106,7 @@ func tracingWrapper(next handlerWrapperFunc) handlerWrapperFunc {
 
 		span.SetAttributes(
 			attribute.String("request_status", status.String()),
-			attribute.String("status_source", string(errorSourceFromContext(ctx))),
+			attribute.String("status_source", string(ErrorSourceFromContext(ctx))),
 		)
 
 		if err != nil {
@@ -136,7 +136,7 @@ func logWrapper(next handlerWrapperFunc) handlerWrapperFunc {
 			logParams = append(logParams, "error", err)
 		}
 
-		logParams = append(logParams, "statusSource", string(errorSourceFromContext(ctx)))
+		logParams = append(logParams, "statusSource", string(ErrorSourceFromContext(ctx)))
 
 		if status > RequestStatusCancelled {
 			logFunc = ctxLogger.Error

backend/adapter_utils_test.go

@@ -19,7 +19,7 @@ func TestErrorWrapper(t *testing.T) {
 		status, err := wrapper(ctx)
 		require.ErrorIs(t, err, actualErr)
 		require.Equal(t, RequestStatusError, status)
-		require.Equal(t, DefaultErrorSource, errorSourceFromContext(ctx))
+		require.Equal(t, DefaultErrorSource, ErrorSourceFromContext(ctx))
 	})
 
 	t.Run("Downstream error should set downstream error source in context", func(t *testing.T) {
@@ -32,6 +32,6 @@ func TestErrorWrapper(t *testing.T) {
 		status, err := wrapper(ctx)
 		require.ErrorIs(t, err, actualErr)
 		require.Equal(t, RequestStatusError, status)
-		require.Equal(t, ErrorSourceDownstream, errorSourceFromContext(ctx))
+		require.Equal(t, ErrorSourceDownstream, ErrorSourceFromContext(ctx))
 	})
 }

backend/data_adapter.go

@@ -46,13 +46,19 @@ func (a *dataSDKAdapter) QueryData(ctx context.Context, req *pluginv2.QueryDataR
 				continue
 			}
 
-			// 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 !r.ErrorSource.IsValid() {
+				// if the error is a downstream error, set error source to downstream, otherwise plugin.
+				if IsDownstreamError(r.Error) {
+					r.ErrorSource = ErrorSourceDownstream
+				} else {
+					r.ErrorSource = ErrorSourcePlugin
+				}
+				resp.Responses[refID] = r
 			}
 
 			if !r.Status.IsValid() {
 				r.Status = statusFromError(r.Error)
+				resp.Responses[refID] = r
 			}
 
 			if r.ErrorSource == ErrorSourceDownstream {

backend/data_adapter_test.go

@@ -224,11 +224,31 @@ func TestQueryData(t *testing.T) {
 					require.NoError(t, err)
 				}
 
-				ss := errorSourceFromContext(actualCtx)
+				ss := ErrorSourceFromContext(actualCtx)
 				require.Equal(t, tc.expErrorSource, ss)
 			})
 		}
 	})
+
+	t.Run("QueryData response without valid error source error should set error source", func(t *testing.T) {
+		someErr := errors.New("oops")
+		downstreamErr := DownstreamError(someErr)
+		a := newDataSDKAdapter(QueryDataHandlerFunc(func(_ context.Context, _ *QueryDataRequest) (*QueryDataResponse, error) {
+			return &QueryDataResponse{
+				Responses: map[string]DataResponse{
+					"A": {Error: someErr},
+					"B": {Error: downstreamErr},
+				},
+			}, nil
+		}))
+		resp, err := a.QueryData(context.Background(), &pluginv2.QueryDataRequest{
+			PluginContext: &pluginv2.PluginContext{},
+		})
+
+		require.NoError(t, err)
+		require.Equal(t, ErrorSourcePlugin, ErrorSource(resp.Responses["A"].ErrorSource))
+		require.Equal(t, ErrorSourceDownstream, ErrorSource(resp.Responses["B"].ErrorSource))
+	})
 }
 
 var finalRoundTripper = httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {

backend/error_source.go

@@ -38,15 +38,20 @@ func IsDownstreamHTTPError(err error) bool {
 	return status.IsDownstreamHTTPError(err)
 }
 
+// DownstreamError creates a new error with status [ErrorSourceDownstream].
 func DownstreamError(err error) error {
 	return status.DownstreamError(err)
 }
 
+// DownstreamErrorf creates a new error with status [ErrorSourceDownstream] and formats
+// according to a format specifier and returns the string as a value that satisfies error.
 func DownstreamErrorf(format string, a ...any) error {
 	return DownstreamError(fmt.Errorf(format, a...))
 }
 
-func errorSourceFromContext(ctx context.Context) ErrorSource {
+// ErrorSourceFromContext returns the error source stored in the context.
+// If no error source is stored in the context, [DefaultErrorSource] is returned.
+func ErrorSourceFromContext(ctx context.Context) ErrorSource {
 	return status.SourceFromContext(ctx)
 }
 

backend/error_source_middleware.go

@@ -0,0 +1,135 @@
+package backend
+
+import (
+	"context"
+	"errors"
+	"fmt"
+)
+
+// NewErrorSourceMiddleware returns a new backend.HandlerMiddleware that sets the error source in the
+// context.Context, based on returned errors or query data response errors.
+// If at least one query data response has a "downstream" error source and there isn't one with a "plugin" error source,
+// the error source in the context is set to "downstream".
+func NewErrorSourceMiddleware() HandlerMiddleware {
+	return HandlerMiddlewareFunc(func(next Handler) Handler {
+		return &ErrorSourceMiddleware{
+			BaseHandler: NewBaseHandler(next),
+		}
+	})
+}
+
+type ErrorSourceMiddleware struct {
+	BaseHandler
+}
+
+func (m *ErrorSourceMiddleware) handleDownstreamError(ctx context.Context, err error) error {
+	if err == nil {
+		return nil
+	}
+
+	if IsDownstreamError(err) {
+		if innerErr := WithDownstreamErrorSource(ctx); innerErr != nil {
+			return fmt.Errorf("failed to set downstream error source: %w", errors.Join(innerErr, err))
+		}
+	}
+
+	return err
+}
+
+func (m *ErrorSourceMiddleware) QueryData(ctx context.Context, req *QueryDataRequest) (*QueryDataResponse, error) {
+	resp, err := m.BaseHandler.QueryData(ctx, req)
+	// we want to always process the error here first before checking anything else
+	// because we want the opportunity to set the error source in the context.
+	err = m.handleDownstreamError(ctx, err)
+
+	// no point in continue if we have an error or no response to process, so we return early.
+	if err != nil || resp == nil || len(resp.Responses) == 0 {
+		return resp, err
+	}
+
+	// Set downstream error source in the context if there's at least one response with downstream error source,
+	// and if there's no plugin error
+	var hasPluginError bool
+	var hasDownstreamError bool
+	for refID, r := range resp.Responses {
+		if r.Error == nil {
+			continue
+		}
+
+		if !r.ErrorSource.IsValid() {
+			// if the error is a downstream error, set error source to downstream, otherwise plugin.
+			if IsDownstreamError(r.Error) {
+				r.ErrorSource = ErrorSourceDownstream
+			} else {
+				r.ErrorSource = ErrorSourcePlugin
+			}
+			resp.Responses[refID] = r
+		}
+
+		if !r.Status.IsValid() {
+			r.Status = statusFromError(r.Error)
+			resp.Responses[refID] = r
+		}
+
+		if r.ErrorSource == ErrorSourceDownstream {
+			hasDownstreamError = true
+		} else {
+			hasPluginError = true
+		}
+	}
+
+	// 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 resp, fmt.Errorf("failed to set downstream status source: %w", err)
+		}
+	}
+
+	return resp, err
+}
+
+func (m *ErrorSourceMiddleware) CallResource(ctx context.Context, req *CallResourceRequest, sender CallResourceResponseSender) error {
+	err := m.BaseHandler.CallResource(ctx, req, sender)
+	return m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) CheckHealth(ctx context.Context, req *CheckHealthRequest) (*CheckHealthResult, error) {
+	resp, err := m.BaseHandler.CheckHealth(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) CollectMetrics(ctx context.Context, req *CollectMetricsRequest) (*CollectMetricsResult, error) {
+	resp, err := m.BaseHandler.CollectMetrics(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) SubscribeStream(ctx context.Context, req *SubscribeStreamRequest) (*SubscribeStreamResponse, error) {
+	resp, err := m.BaseHandler.SubscribeStream(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) PublishStream(ctx context.Context, req *PublishStreamRequest) (*PublishStreamResponse, error) {
+	resp, err := m.BaseHandler.PublishStream(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) RunStream(ctx context.Context, req *RunStreamRequest, sender *StreamSender) error {
+	err := m.BaseHandler.RunStream(ctx, req, sender)
+	return m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) ValidateAdmission(ctx context.Context, req *AdmissionRequest) (*ValidationResponse, error) {
+	resp, err := m.BaseHandler.ValidateAdmission(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) MutateAdmission(ctx context.Context, req *AdmissionRequest) (*MutationResponse, error) {
+	resp, err := m.BaseHandler.MutateAdmission(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}
+
+func (m *ErrorSourceMiddleware) ConvertObjects(ctx context.Context, req *ConversionRequest) (*ConversionResponse, error) {
+	resp, err := m.BaseHandler.ConvertObjects(ctx, req)
+	return resp, m.handleDownstreamError(ctx, err)
+}