SSE: Make sure to forward headers, user and cookies/OAuth token (#58897)

Fixes #58793 and Fixes https://github.com/grafana/azure-data-explorer-datasource/issues/513

Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>

pkg/expr/graph.go

@@ -144,11 +144,12 @@ func (s *Service) buildGraph(req *Request) (*simple.DirectedGraph, error) {
 		}
 
 		rn := &rawNode{
-			Query:      rawQueryProp,
-			RefID:      query.RefID,
-			TimeRange:  query.TimeRange,
-			QueryType:  query.QueryType,
-			DataSource: query.DataSource,
+			Query:         rawQueryProp,
+			RefID:         query.RefID,
+			TimeRange:     query.TimeRange,
+			QueryType:     query.QueryType,
+			DataSource:    query.DataSource,
+			QueryEnricher: query.QueryEnricher,
 		}
 
 		var node Node

pkg/expr/nodes.go

@@ -42,11 +42,12 @@ type baseNode struct {
 }
 
 type rawNode struct {
-	RefID      string `json:"refId"`
-	Query      map[string]interface{}
-	QueryType  string
-	TimeRange  TimeRange
-	DataSource *datasources.DataSource
+	RefID         string `json:"refId"`
+	Query         map[string]interface{}
+	QueryType     string
+	TimeRange     TimeRange
+	DataSource    *datasources.DataSource
+	QueryEnricher QueryDataRequestEnricher
 }
 
 func (rn *rawNode) GetCommandType() (c CommandType, err error) {
@@ -139,8 +140,9 @@ const (
 // DSNode is a DPNode that holds a datasource request.
 type DSNode struct {
 	baseNode
-	query      json.RawMessage
-	datasource *datasources.DataSource
+	query         json.RawMessage
+	datasource    *datasources.DataSource
+	queryEnricher QueryDataRequestEnricher
 
 	orgID      int64
 	queryType  string
@@ -169,14 +171,15 @@ func (s *Service) buildDSNode(dp *simple.DirectedGraph, rn *rawNode, req *Reques
 			id:    dp.NewNode().ID(),
 			refID: rn.RefID,
 		},
-		orgID:      req.OrgId,
-		query:      json.RawMessage(encodedQuery),
-		queryType:  rn.QueryType,
-		intervalMS: defaultIntervalMS,
-		maxDP:      defaultMaxDP,
-		timeRange:  rn.TimeRange,
-		request:    *req,
-		datasource: rn.DataSource,
+		orgID:         req.OrgId,
+		query:         json.RawMessage(encodedQuery),
+		queryType:     rn.QueryType,
+		intervalMS:    defaultIntervalMS,
+		maxDP:         defaultMaxDP,
+		timeRange:     rn.TimeRange,
+		request:       *req,
+		datasource:    rn.DataSource,
+		queryEnricher: rn.QueryEnricher,
 	}
 
 	var floatIntervalMS float64
@@ -211,24 +214,29 @@ func (dn *DSNode) Execute(ctx context.Context, now time.Time, _ mathexp.Vars, s
 		OrgID:                      dn.orgID,
 		DataSourceInstanceSettings: dsInstanceSettings,
 		PluginID:                   dn.datasource.Type,
+		User:                       dn.request.User,
 	}
 
-	q := []backend.DataQuery{
-		{
-			RefID:         dn.refID,
-			MaxDataPoints: dn.maxDP,
-			Interval:      time.Duration(int64(time.Millisecond) * dn.intervalMS),
-			JSON:          dn.query,
-			TimeRange:     dn.timeRange.AbsoluteTime(now),
-			QueryType:     dn.queryType,
-		},
-	}
-
-	resp, err := s.dataService.QueryData(ctx, &backend.QueryDataRequest{
+	req := &backend.QueryDataRequest{
 		PluginContext: pc,
-		Queries:       q,
-		Headers:       dn.request.Headers,
-	})
+		Queries: []backend.DataQuery{
+			{
+				RefID:         dn.refID,
+				MaxDataPoints: dn.maxDP,
+				Interval:      time.Duration(int64(time.Millisecond) * dn.intervalMS),
+				JSON:          dn.query,
+				TimeRange:     dn.timeRange.AbsoluteTime(now),
+				QueryType:     dn.queryType,
+			},
+		},
+		Headers: dn.request.Headers,
+	}
+
+	if dn.queryEnricher != nil {
+		ctx = dn.queryEnricher(ctx, req)
+	}
+
+	resp, err := s.dataService.QueryData(ctx, req)
 	if err != nil {
 		return mathexp.Results{}, err
 	}

pkg/expr/transform.go

@@ -35,14 +35,19 @@ type Request struct {
 	Debug   bool
 	OrgId   int64
 	Queries []Query
+	User    *backend.User
 }
 
+// QueryDataRequestEnricher function definition for enriching a backend.QueryDataRequest request.
+type QueryDataRequestEnricher func(ctx context.Context, req *backend.QueryDataRequest) context.Context
+
 // Query is like plugins.DataSubQuery, but with a a time range, and only the UID
 // for the data source. Also interval is a time.Duration.
 type Query struct {
 	RefID         string
 	TimeRange     TimeRange
 	DataSource    *datasources.DataSource `json:"datasource"`
+	QueryEnricher QueryDataRequestEnricher
 	JSON          json.RawMessage
 	Interval      time.Duration
 	QueryType     string