From 48a1dae834b9567f33ae6b7b2bc1e74b741bb868 Mon Sep 17 00:00:00 2001
From: Giuseppe Guerra <giuseppe.guerra@grafana.com>
Date: Thu, 19 Oct 2023 11:52:50 +0200
Subject: [PATCH] Plugins: Add contextual logger to streaming methods in
 ContextualLoggerMiddleware (#76761)

---
 .../clientmiddleware/contextual_logger_middleware.go  |  3 +++
 .../pluginsintegration/clientmiddleware/utils.go      | 11 +++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/pkg/services/pluginsintegration/clientmiddleware/contextual_logger_middleware.go b/pkg/services/pluginsintegration/clientmiddleware/contextual_logger_middleware.go
index cf97590294b..029c1968d84 100644
--- a/pkg/services/pluginsintegration/clientmiddleware/contextual_logger_middleware.go
+++ b/pkg/services/pluginsintegration/clientmiddleware/contextual_logger_middleware.go
@@ -57,13 +57,16 @@ func (m *ContextualLoggerMiddleware) CollectMetrics(ctx context.Context, req *ba
 }
 
 func (m *ContextualLoggerMiddleware) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
+	ctx = instrumentContext(ctx, endpointSubscribeStream, req.PluginContext)
 	return m.next.SubscribeStream(ctx, req)
 }
 
 func (m *ContextualLoggerMiddleware) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
+	ctx = instrumentContext(ctx, endpointPublishStream, req.PluginContext)
 	return m.next.PublishStream(ctx, req)
 }
 
 func (m *ContextualLoggerMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
+	ctx = instrumentContext(ctx, endpointRunStream, req.PluginContext)
 	return m.next.RunStream(ctx, req, sender)
 }
diff --git a/pkg/services/pluginsintegration/clientmiddleware/utils.go b/pkg/services/pluginsintegration/clientmiddleware/utils.go
index 4201c3fe417..462d4250f1a 100644
--- a/pkg/services/pluginsintegration/clientmiddleware/utils.go
+++ b/pkg/services/pluginsintegration/clientmiddleware/utils.go
@@ -9,10 +9,13 @@ const (
 	statusError     = "error"
 	statusCancelled = "cancelled"
 
-	endpointCallResource   = "callResource"
-	endpointCheckHealth    = "checkHealth"
-	endpointCollectMetrics = "collectMetrics"
-	endpointQueryData      = "queryData"
+	endpointCallResource    = "callResource"
+	endpointCheckHealth     = "checkHealth"
+	endpointCollectMetrics  = "collectMetrics"
+	endpointQueryData       = "queryData"
+	endpointSubscribeStream = "subscribeStream"
+	endpointPublishStream   = "publishStream"
+	endpointRunStream       = "runStream"
 )
 
 type callResourceResponseSenderFunc func(res *backend.CallResourceResponse) error