Plugins: Fix status_source always being "plugin" in plugin request logs (#77433)

* Plugins: Fix status_source always being "plugin" in plugin logs

* add tests

* Fix TestInstrumentationMiddlewareStatusSource

pkg/services/pluginsintegration/clientmiddleware/metrics_middleware.go

@@ -3,7 +3,6 @@ package clientmiddleware
 import (
 	"context"
 	"errors"
-	"fmt"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -160,9 +159,6 @@ func (m *MetricsMiddleware) instrumentPluginRequest(ctx context.Context, pluginC
 }
 
 func (m *MetricsMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
-	// Setup plugin request status source
-	ctx = pluginrequestmeta.WithStatusSource(ctx, pluginrequestmeta.StatusSourcePlugin)
-
 	var requestSize float64
 	for _, v := range req.Queries {
 		requestSize += float64(len(v.JSON))
@@ -173,33 +169,7 @@ func (m *MetricsMiddleware) QueryData(ctx context.Context, req *backend.QueryDat
 	var resp *backend.QueryDataResponse
 	err := m.instrumentPluginRequest(ctx, req.PluginContext, endpointQueryData, func(ctx context.Context) (innerErr error) {
 		resp, innerErr = m.next.QueryData(ctx, req)
-		if resp == nil || resp.Responses == nil || !m.features.IsEnabled(featuremgmt.FlagPluginsInstrumentationStatusSource) {
+		return
-			return innerErr
-		}
-
-		// 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
-		for _, r := range resp.Responses {
-			if r.Error == nil {
-				continue
-			}
-			if r.ErrorSource == backend.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 := pluginrequestmeta.WithDownstreamStatusSource(ctx); err != nil {
-				return fmt.Errorf("failed to set downstream status source: %w", err)
-			}
-		}
-
-		return innerErr
 	})
 	return resp, err
 }

pkg/services/pluginsintegration/clientmiddleware/metrics_middleware_test.go

@@ -156,10 +156,12 @@ func TestInstrumentationMiddlewareStatusSource(t *testing.T) {
 	features := featuremgmt.WithFeatures(featuremgmt.FlagPluginsInstrumentationStatusSource)
 	metricsMw := newMetricsMiddleware(promRegistry, pluginsRegistry, features)
 	cdt := clienttest.NewClientDecoratorTest(t, clienttest.WithMiddlewares(
+		NewPluginRequestMetaMiddleware(),
 		plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
 			metricsMw.next = next
 			return metricsMw
 		}),
+		NewStatusSourceMiddleware(),
 	))
 
 	t.Run("Metrics", func(t *testing.T) {

pkg/services/pluginsintegration/clientmiddleware/plugin_request_meta_middleware.go

@@ -0,0 +1,69 @@
+package clientmiddleware
+
+import (
+	"context"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+
+	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/plugins/pluginrequestmeta"
+)
+
+// NewPluginRequestMetaMiddleware returns a new plugins.ClientMiddleware that sets up the default
+// values for the plugin request meta in the context.Context. All middlewares that are executed
+// after this one are be able to access plugin request meta via the pluginrequestmeta package.
+func NewPluginRequestMetaMiddleware() plugins.ClientMiddleware {
+	return plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
+		return &PluginRequestMetaMiddleware{
+			next:                next,
+			defaultStatusSource: pluginrequestmeta.StatusSourcePlugin,
+		}
+	})
+}
+
+type PluginRequestMetaMiddleware struct {
+	next                plugins.Client
+	defaultStatusSource pluginrequestmeta.StatusSource
+}
+
+func (m *PluginRequestMetaMiddleware) withDefaultPluginRequestMeta(ctx context.Context) context.Context {
+	// Setup plugin request status source
+	ctx = pluginrequestmeta.WithStatusSource(ctx, m.defaultStatusSource)
+
+	return ctx
+}
+
+func (m *PluginRequestMetaMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.QueryData(ctx, req)
+}
+
+func (m *PluginRequestMetaMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.CallResource(ctx, req, sender)
+}
+
+func (m *PluginRequestMetaMiddleware) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.CheckHealth(ctx, req)
+}
+
+func (m *PluginRequestMetaMiddleware) CollectMetrics(ctx context.Context, req *backend.CollectMetricsRequest) (*backend.CollectMetricsResult, error) {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.CollectMetrics(ctx, req)
+}
+
+func (m *PluginRequestMetaMiddleware) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.SubscribeStream(ctx, req)
+}
+
+func (m *PluginRequestMetaMiddleware) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.PublishStream(ctx, req)
+}
+
+func (m *PluginRequestMetaMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
+	ctx = m.withDefaultPluginRequestMeta(ctx)
+	return m.next.RunStream(ctx, req, sender)
+}

pkg/services/pluginsintegration/clientmiddleware/plugin_request_meta_middleware_test.go

@@ -0,0 +1,40 @@
+package clientmiddleware
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/plugins/manager/client/clienttest"
+	"github.com/grafana/grafana/pkg/plugins/pluginrequestmeta"
+)
+
+func TestPluginRequestMetaMiddleware(t *testing.T) {
+	t.Run("default", func(t *testing.T) {
+		cdt := clienttest.NewClientDecoratorTest(t,
+			clienttest.WithMiddlewares(NewPluginRequestMetaMiddleware()),
+		)
+		_, err := cdt.Decorator.QueryData(context.Background(), &backend.QueryDataRequest{})
+		require.NoError(t, err)
+		ss := pluginrequestmeta.StatusSourceFromContext(cdt.QueryDataCtx)
+		require.Equal(t, pluginrequestmeta.StatusSourcePlugin, ss)
+	})
+
+	t.Run("other value", func(t *testing.T) {
+		cdt := clienttest.NewClientDecoratorTest(t,
+			clienttest.WithMiddlewares(plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
+				return &PluginRequestMetaMiddleware{
+					next:                next,
+					defaultStatusSource: "test",
+				}
+			})),
+		)
+		_, err := cdt.Decorator.QueryData(context.Background(), &backend.QueryDataRequest{})
+		require.NoError(t, err)
+		ss := pluginrequestmeta.StatusSourceFromContext(cdt.QueryDataCtx)
+		require.Equal(t, pluginrequestmeta.StatusSource("test"), ss)
+	})
+}

pkg/services/pluginsintegration/clientmiddleware/status_source_middleware.go

@@ -0,0 +1,83 @@
+package clientmiddleware
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+
+	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/plugins/pluginrequestmeta"
+)
+
+// NewStatusSourceMiddleware returns a new plugins.ClientMiddleware that sets the status source in the
+// plugin request meta stored in the context.Context, according to the query data responses returned by QueryError.
+// If at least one query data response has a "downstream" status source and there isn't one with a "plugin" status source,
+// the plugin request meta in the context is set to "downstream".
+func NewStatusSourceMiddleware() plugins.ClientMiddleware {
+	return plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
+		return &StatusSourceMiddleware{
+			next: next,
+		}
+	})
+}
+
+type StatusSourceMiddleware struct {
+	next plugins.Client
+}
+
+func (m *StatusSourceMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	resp, err := m.next.QueryData(ctx, req)
+	if resp == nil || len(resp.Responses) == 0 {
+		return resp, err
+	}
+
+	// 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
+	for _, r := range resp.Responses {
+		if r.Error == nil {
+			continue
+		}
+		if r.ErrorSource == backend.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 := pluginrequestmeta.WithDownstreamStatusSource(ctx); err != nil {
+			return resp, fmt.Errorf("failed to set downstream status source: %w", err)
+		}
+	}
+
+	return resp, err
+}
+
+func (m *StatusSourceMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	return m.next.CallResource(ctx, req, sender)
+}
+
+func (m *StatusSourceMiddleware) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	return m.next.CheckHealth(ctx, req)
+}
+
+func (m *StatusSourceMiddleware) CollectMetrics(ctx context.Context, req *backend.CollectMetricsRequest) (*backend.CollectMetricsResult, error) {
+	return m.next.CollectMetrics(ctx, req)
+}
+
+func (m *StatusSourceMiddleware) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
+	return m.next.SubscribeStream(ctx, req)
+}
+
+func (m *StatusSourceMiddleware) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
+	return m.next.PublishStream(ctx, req)
+}
+
+func (m *StatusSourceMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
+	return m.next.RunStream(ctx, req, sender)
+}

pkg/services/pluginsintegration/clientmiddleware/status_source_middleware_test.go

@@ -0,0 +1,88 @@
+package clientmiddleware
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/plugins/manager/client/clienttest"
+	"github.com/grafana/grafana/pkg/plugins/pluginrequestmeta"
+)
+
+func TestStatusSourceMiddleware(t *testing.T) {
+	someErr := errors.New("oops")
+
+	for _, tc := range []struct {
+		name string
+
+		queryDataResponse *backend.QueryDataResponse
+
+		expStatusSource pluginrequestmeta.StatusSource
+	}{
+		{
+			name:              `no error should be "plugin" status source`,
+			queryDataResponse: nil,
+			expStatusSource:   pluginrequestmeta.StatusSourcePlugin,
+		},
+		{
+			name: `single downstream error should be "downstream" status source`,
+			queryDataResponse: &backend.QueryDataResponse{
+				Responses: map[string]backend.DataResponse{
+					"A": {Error: someErr, ErrorSource: backend.ErrorSourceDownstream},
+				},
+			},
+			expStatusSource: pluginrequestmeta.StatusSourceDownstream,
+		},
+		{
+			name: `single plugin error should be "plugin" status source`,
+			queryDataResponse: &backend.QueryDataResponse{
+				Responses: map[string]backend.DataResponse{
+					"A": {Error: someErr, ErrorSource: backend.ErrorSourcePlugin},
+				},
+			},
+			expStatusSource: pluginrequestmeta.StatusSourcePlugin,
+		},
+		{
+			name: `multiple downstream errors should be "downstream" status source`,
+			queryDataResponse: &backend.QueryDataResponse{
+				Responses: map[string]backend.DataResponse{
+					"A": {Error: someErr, ErrorSource: backend.ErrorSourceDownstream},
+					"B": {Error: someErr, ErrorSource: backend.ErrorSourceDownstream},
+				},
+			},
+			expStatusSource: pluginrequestmeta.StatusSourceDownstream,
+		},
+		{
+			name: `single plugin error mixed with downstream errors should be "plugin" status source`,
+			queryDataResponse: &backend.QueryDataResponse{
+				Responses: map[string]backend.DataResponse{
+					"A": {Error: someErr, ErrorSource: backend.ErrorSourceDownstream},
+					"B": {Error: someErr, ErrorSource: backend.ErrorSourcePlugin},
+					"C": {Error: someErr, ErrorSource: backend.ErrorSourceDownstream},
+				},
+			},
+			expStatusSource: pluginrequestmeta.StatusSourcePlugin,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			cdt := clienttest.NewClientDecoratorTest(t,
+				clienttest.WithMiddlewares(
+					NewPluginRequestMetaMiddleware(),
+					NewStatusSourceMiddleware(),
+				),
+			)
+			cdt.TestClient.QueryDataFunc = func(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+				cdt.QueryDataCtx = ctx
+				return tc.queryDataResponse, nil
+			}
+
+			_, _ = cdt.Decorator.QueryData(context.Background(), &backend.QueryDataRequest{})
+
+			ss := pluginrequestmeta.StatusSourceFromContext(cdt.QueryDataCtx)
+			require.Equal(t, tc.expStatusSource, ss)
+		})
+	}
+}

pkg/services/pluginsintegration/pluginsintegration.go

@@ -150,8 +150,16 @@ func NewClientDecorator(
 }
 
 func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthTokenService, tracer tracing.Tracer, cachingService caching.CachingService, features *featuremgmt.FeatureManager, promRegisterer prometheus.Registerer, registry registry.Service) []plugins.ClientMiddleware {
+	var middlewares []plugins.ClientMiddleware
+
+	if features.IsEnabled(featuremgmt.FlagPluginsInstrumentationStatusSource) {
+		middlewares = []plugins.ClientMiddleware{
+			clientmiddleware.NewPluginRequestMetaMiddleware(),
+		}
+	}
+
 	skipCookiesNames := []string{cfg.LoginCookieName}
-	middlewares := []plugins.ClientMiddleware{
+	middlewares = append(middlewares,
 		clientmiddleware.NewTracingMiddleware(tracer),
 		clientmiddleware.NewMetricsMiddleware(promRegisterer, registry, features),
 		clientmiddleware.NewContextualLoggerMiddleware(),
@@ -161,7 +169,7 @@ func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthToken
 		clientmiddleware.NewOAuthTokenMiddleware(oAuthTokenService),
 		clientmiddleware.NewCookiesMiddleware(skipCookiesNames),
 		clientmiddleware.NewResourceResponseMiddleware(),
-	}
+	)
 
 	// Placing the new service implementation behind a feature flag until it is known to be stable
 	if features.IsEnabled(featuremgmt.FlagUseCachingService) {
@@ -178,5 +186,11 @@ func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthToken
 
 	middlewares = append(middlewares, clientmiddleware.NewHTTPClientMiddleware())
 
+	if features.IsEnabled(featuremgmt.FlagPluginsInstrumentationStatusSource) {
+		// StatusSourceMiddleware should be at the very bottom, or any middlewares below it won't see the
+		// correct status source in their context.Context
+		middlewares = append(middlewares, clientmiddleware.NewStatusSourceMiddleware())
+	}
+
 	return middlewares
 }