Return appropriate status type plugin request status handler func (#1036

)

* check for inner data resp err

* update test

* update

* fixup

* check for cancelled

* fix err messages

* update cancelled err check

backend/adapter_utils.go

@@ -135,7 +135,7 @@ func logWrapper(next handlerWrapperFunc) handlerWrapperFunc {
 
 		ctxLogger := Logger.FromContext(ctx)
 		logFunc := ctxLogger.Debug
-		if status > RequestStatusOK {
+		if status > RequestStatusCancelled {
 			logFunc = ctxLogger.Error
 		}
 

backend/data_adapter.go

@@ -33,30 +33,58 @@ func (a *dataSDKAdapter) QueryData(ctx context.Context, req *pluginv2.QueryDataR
 			return RequestStatusFromError(innerErr), innerErr
 		}
 
+		if isCancelledError(innerErr) {
+			return RequestStatusCancelled, nil
+		}
+
 		// 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
 		for _, r := range resp.Responses {
 			if r.Error == nil {
 				continue
 			}
+
+			if isCancelledError(r.Error) {
+				hasCancelledError = true
+			}
 			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
+		}
+
 		// 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
 		}
 
-		return RequestStatusFromError(innerErr), innerErr
+		return RequestStatusOK, nil
 	})
 	if err != nil {
 		return nil, err

backend/request_status.go

@@ -32,7 +32,7 @@ func RequestStatusFromError(err error) RequestStatus {
 	status := RequestStatusOK
 	if err != nil {
 		status = RequestStatusError
-		if errors.Is(err, context.Canceled) || grpcstatus.Code(err) == grpccodes.Canceled {
+		if isCancelledError(err) {
 			status = RequestStatusCancelled
 		}
 	}
@@ -101,3 +101,7 @@ func RequestStatusFromProtoQueryDataResponse(res *pluginv2.QueryDataResponse, er
 
 	return status
 }
+
+func isCancelledError(err error) bool {
+	return errors.Is(err, context.Canceled) || grpcstatus.Code(err) == grpccodes.Canceled
+}