Logging: Add WithContextualAttributes to pass log params based on the given context

pkg/infra/log/log.go

@@ -51,6 +51,14 @@ func init() {
 	}
 	logger := level.NewFilter(format(os.Stderr), level.AllowInfo())
 	root = newManager(logger)
+
+	RegisterContextualLogProvider(func(ctx context.Context) ([]any, bool) {
+		pFromCtx := ctx.Value(logParamsContextKey{})
+		if pFromCtx != nil {
+			return pFromCtx.([]any), true
+		}
+		return nil, false
+	})
 }
 
 // logManager manage loggers
@@ -270,6 +278,18 @@ func RegisterContextualLogProvider(mw ContextualLogProviderFunc) {
 	ctxLogProviders = append(ctxLogProviders, mw)
 }
 
+type logParamsContextKey struct{}
+
+// WithContextualAttributes adds contextual attributes to the logger based on the given context.
+// That allows loggers further down the chain to automatically log those attributes.
+func WithContextualAttributes(ctx context.Context, logParams []any) context.Context {
+	p := logParams
+	if ctx.Value(logParamsContextKey{}) != nil {
+		p = append(ctx.Value(logParamsContextKey{}).([]any), logParams...)
+	}
+	return context.WithValue(ctx, logParamsContextKey{}, p)
+}
+
 var logLevels = map[string]level.Option{
 	"trace":    level.AllowDebug(),
 	"debug":    level.AllowDebug(),

pkg/infra/log/log_test.go

@@ -267,6 +267,84 @@ func TestGetFilters(t *testing.T) {
 	})
 }
 
+func TestWithContextualAttributes_appendsContext(t *testing.T) {
+	t.Run("Logs arguments from context", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		// logs `"k1", "v1"` with the first context
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+	})
+	t.Run("Does not log arguments from different context", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		// logs `"k1", "v1"` with the first context
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+		// does not log `"k1", "v1"` with the new context
+		ctx = context.Background()
+		ls = New("test").FromContext(ctx)
+
+		ls.Info("hello", "k2", "v2")
+
+		require.Len(t, scenario.loggedArgs, 2)
+		scenario.ValidateLineEquality(t, 1, []any{
+			"logger", "test",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k2", "v2",
+		})
+	})
+	t.Run("Appends arguments set previously", func(t *testing.T) {
+		scenario := newLoggerScenario(t, false)
+
+		ctx := context.Background()
+		ctx = WithContextualAttributes(ctx, []any{"k1", "v1"})
+		ctx = WithContextualAttributes(ctx, []any{"k2", "v2"})
+		ls := New("test").FromContext(ctx)
+
+		ls.Info("hello", "k3", "v3")
+
+		require.Len(t, scenario.loggedArgs, 1)
+		scenario.ValidateLineEquality(t, 0, []any{
+			"logger", "test",
+			"k1", "v1",
+			"k2", "v2",
+			"t", scenario.mockedTime,
+			level.Key(), level.InfoValue(),
+			"msg", "hello",
+			"k3", "v3",
+		})
+	})
+}
+
 type scenarioContext struct {
 	loggedArgs [][]any
 	mockedTime time.Time
@@ -287,7 +365,12 @@ func (s *scenarioContext) ValidateLineEquality(t testing.TB, n int, expected []a
 	}
 }
 
-func newLoggerScenario(t testing.TB) *scenarioContext {
+func newLoggerScenario(t testing.TB, resetCtxLogProviders ...bool) *scenarioContext {
+	clearProviders := true
+	if len(resetCtxLogProviders) > 0 {
+		clearProviders = resetCtxLogProviders[0]
+	}
+
 	t.Helper()
 
 	scenario := &scenarioContext{
@@ -308,11 +391,13 @@ func newLoggerScenario(t testing.TB) *scenarioContext {
 		now = origNow
 	})
 
-	origContextHandlers := ctxLogProviders
-	ctxLogProviders = []ContextualLogProviderFunc{}
-	t.Cleanup(func() {
-		ctxLogProviders = origContextHandlers
-	})
+	if clearProviders {
+		origContextHandlers := ctxLogProviders
+		ctxLogProviders = []ContextualLogProviderFunc{}
+		t.Cleanup(func() {
+			ctxLogProviders = origContextHandlers
+		})
+	}
 
 	root = newManager(l)
 	return scenario

pkg/plugins/backendplugin/instrumentation/instrumentation.go

@@ -61,13 +61,14 @@ const (
 var logger = plog.New("plugin.instrumentation")
 
 // instrumentPluginRequest instruments success rate and latency of `fn`
-func instrumentPluginRequest(ctx context.Context, cfg Cfg, pluginCtx *backend.PluginContext, endpoint string, fn func() error) error {
+func instrumentPluginRequest(ctx context.Context, cfg Cfg, pluginCtx *backend.PluginContext, endpoint string, fn func(ctx context.Context) error) error {
 	status := statusOK
 
 	start := time.Now()
 	timeBeforePluginRequest := log.TimeSinceStart(ctx, start)
 
-	err := fn()
+	ctx = instrumentContext(ctx, endpoint, *pluginCtx)
+	err := fn(ctx)
 	if err != nil {
 		status = statusError
 		if errors.Is(err, context.Canceled) {
@@ -129,31 +130,40 @@ func instrumentPluginRequest(ctx context.Context, cfg Cfg, pluginCtx *backend.Pl
 	return err
 }
 
+func instrumentContext(ctx context.Context, endpoint string, pCtx backend.PluginContext) context.Context {
+	p := []any{"endpoint", endpoint, "pluginId", pCtx.PluginID}
+	if pCtx.DataSourceInstanceSettings != nil {
+		p = append(p, "dsName", pCtx.DataSourceInstanceSettings.Name)
+		p = append(p, "dsUID", pCtx.DataSourceInstanceSettings.UID)
+	}
+	return log.WithContextualAttributes(ctx, p)
+}
+
 type Cfg struct {
 	LogDatasourceRequests bool
 	Target                backendplugin.Target
 }
 
 // InstrumentCollectMetrics instruments collectMetrics.
-func InstrumentCollectMetrics(ctx context.Context, req *backend.PluginContext, cfg Cfg, fn func() error) error {
+func InstrumentCollectMetrics(ctx context.Context, req *backend.PluginContext, cfg Cfg, fn func(ctx context.Context) error) error {
 	return instrumentPluginRequest(ctx, cfg, req, endpointCollectMetrics, fn)
 }
 
 // InstrumentCheckHealthRequest instruments checkHealth.
-func InstrumentCheckHealthRequest(ctx context.Context, req *backend.PluginContext, cfg Cfg, fn func() error) error {
+func InstrumentCheckHealthRequest(ctx context.Context, req *backend.PluginContext, cfg Cfg, fn func(ctx context.Context) error) error {
 	return instrumentPluginRequest(ctx, cfg, req, endpointCheckHealth, fn)
 }
 
 // InstrumentCallResourceRequest instruments callResource.
-func InstrumentCallResourceRequest(ctx context.Context, req *backend.PluginContext, cfg Cfg, requestSize float64, fn func() error) error {
+func InstrumentCallResourceRequest(ctx context.Context, req *backend.PluginContext, cfg Cfg, requestSize float64, fn func(ctx context.Context) error) error {
 	pluginRequestSizeHistogram.WithLabelValues("grafana-backend", req.PluginID, endpointCallResource,
 		string(cfg.Target)).Observe(requestSize)
 	return instrumentPluginRequest(ctx, cfg, req, endpointCallResource, fn)
 }
 
 // InstrumentQueryDataRequest instruments success rate and latency of query data requests.
 func InstrumentQueryDataRequest(ctx context.Context, req *backend.PluginContext, cfg Cfg,
-	requestSize float64, fn func() error) error {
+	requestSize float64, fn func(ctx context.Context) error) error {
 	pluginRequestSizeHistogram.WithLabelValues("grafana-backend", req.PluginID, endpointQueryData,
 		string(cfg.Target)).Observe(requestSize)
 	return instrumentPluginRequest(ctx, cfg, req, endpointQueryData, fn)

pkg/plugins/manager/client/client.go

@@ -61,7 +61,7 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest)
 	err := instrumentation.InstrumentQueryDataRequest(ctx, &req.PluginContext, instrumentation.Cfg{
 		LogDatasourceRequests: s.cfg.LogDatasourceRequests,
 		Target:                p.Target(),
-	}, totalBytes, func() (innerErr error) {
+	}, totalBytes, func(ctx context.Context) (innerErr error) {
 		resp, innerErr = p.QueryData(ctx, req)
 		return
 	})
@@ -108,7 +108,7 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq
 	err := instrumentation.InstrumentCallResourceRequest(ctx, &req.PluginContext, instrumentation.Cfg{
 		LogDatasourceRequests: s.cfg.LogDatasourceRequests,
 		Target:                p.Target(),
-	}, totalBytes, func() error {
+	}, totalBytes, func(ctx context.Context) error {
 		removeConnectionHeaders(req.Headers)
 		removeHopByHopHeaders(req.Headers)
 		removeNonAllowedHeaders(req.Headers)
@@ -157,7 +157,7 @@ func (s *Service) CollectMetrics(ctx context.Context, req *backend.CollectMetric
 	err := instrumentation.InstrumentCollectMetrics(ctx, &req.PluginContext, instrumentation.Cfg{
 		LogDatasourceRequests: s.cfg.LogDatasourceRequests,
 		Target:                p.Target(),
-	}, func() (innerErr error) {
+	}, func(ctx context.Context) (innerErr error) {
 		resp, innerErr = p.CollectMetrics(ctx, req)
 		return
 	})
@@ -182,7 +182,7 @@ func (s *Service) CheckHealth(ctx context.Context, req *backend.CheckHealthReque
 	err := instrumentation.InstrumentCheckHealthRequest(ctx, &req.PluginContext, instrumentation.Cfg{
 		LogDatasourceRequests: s.cfg.LogDatasourceRequests,
 		Target:                p.Target(),
-	}, func() (innerErr error) {
+	}, func(ctx context.Context) (innerErr error) {
 		resp, innerErr = p.CheckHealth(ctx, req)
 		return
 	})

pkg/plugins/manager/client/decorator.go

@@ -32,6 +32,7 @@ func (d *Decorator) QueryData(ctx context.Context, req *backend.QueryDataRequest
 	}
 
 	client := clientFromMiddlewares(d.middlewares, d.client)
+
 	return client.QueryData(ctx, req)
 }