feat: add normalizedQuery to query plan and request info to trace

wundergraph · Jan 25, 2025 · 4c4a0f0 · 4c4a0f0
1 parent 394a730
commit 4c4a0f0
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 13 deletions.
diff --git a/v2/pkg/engine/resolve/context.go b/v2/pkg/engine/resolve/context.go
@@ -193,6 +193,8 @@ type PhaseStats struct {
 	DurationSinceStartPretty string `json:"duration_since_start_pretty"`
 }
 
+type requestContextKey struct{}
+
 func SetTraceStart(ctx context.Context, predictableDebugTimings bool) context.Context {
 	info := &TraceInfo{}
 	if predictableDebugTimings {
@@ -267,3 +269,16 @@ func SetPlannerStats(ctx context.Context, stats PhaseStats) {
 	}
 	info.PlannerStats = SetDebugStats(info, stats, 4)
 }
+
+func GetRequest(ctx context.Context) *RequestData {
+	// The context might not have trace info, in that case we return nil
+	req, ok := ctx.Value(requestContextKey{}).(*RequestData)
+	if !ok {
+		return nil
+	}
+	return req
+}
+
+func SetRequest(ctx context.Context, r *RequestData) context.Context {
+	return context.WithValue(ctx, requestContextKey{}, r)
+}
diff --git a/v2/pkg/engine/resolve/fetchtree.go b/v2/pkg/engine/resolve/fetchtree.go
@@ -8,9 +8,10 @@ import (
 type FetchTreeNode struct {
 	Kind FetchTreeNodeKind `json:"kind"`
 	// Only set for subscription
-	Trigger    *FetchTreeNode   `json:"trigger"`
-	Item       *FetchItem       `json:"item"`
-	ChildNodes []*FetchTreeNode `json:"child_nodes"`
+	Trigger         *FetchTreeNode   `json:"trigger"`
+	Item            *FetchItem       `json:"item"`
+	ChildNodes      []*FetchTreeNode `json:"child_nodes"`
+	NormalizedQuery string           `json:"normalized_query"`
 }
 
 type FetchTreeNodeKind string
@@ -147,11 +148,12 @@ func (n *FetchTreeNode) Trace() *FetchTreeTraceNode {
 }
 
 type FetchTreeQueryPlanNode struct {
-	Version  string                    `json:"version,omitempty"`
-	Kind     FetchTreeNodeKind         `json:"kind"`
-	Trigger  *FetchTreeQueryPlan       `json:"trigger,omitempty"`
-	Children []*FetchTreeQueryPlanNode `json:"children,omitempty"`
-	Fetch    *FetchTreeQueryPlan       `json:"fetch,omitempty"`
+	Version         string                    `json:"version,omitempty"`
+	Kind            FetchTreeNodeKind         `json:"kind"`
+	Trigger         *FetchTreeQueryPlan       `json:"trigger,omitempty"`
+	Children        []*FetchTreeQueryPlanNode `json:"children,omitempty"`
+	Fetch           *FetchTreeQueryPlan       `json:"fetch,omitempty"`
+	NormalizedQuery string                    `json:"normalizedQuery,omitempty"`
 }
 
 type FetchTreeQueryPlan struct {
@@ -194,7 +196,8 @@ func (n *FetchTreeNode) queryPlan() *FetchTreeQueryPlanNode {
 		return nil
 	}
 	queryPlan := &FetchTreeQueryPlanNode{
-		Kind: n.Kind,
+		Kind:            n.Kind,
+		NormalizedQuery: n.NormalizedQuery,
 	}
 	switch n.Kind {
 	case FetchTreeNodeKindSingle:

diff --git a/v2/pkg/engine/resolve/trace.go b/v2/pkg/engine/resolve/trace.go
@@ -2,6 +2,8 @@ package resolve
 
 import (
 	"context"
+	"encoding/json"
+	"net/http"
 )
 
 type TraceOptions struct {
@@ -59,16 +61,30 @@ func (r *TraceOptions) DisableAll() {
 	r.IncludeTraceOutputInResponseExtensions = false
 }
 
+type RequestData struct {
+	Method  string          `json:"method"`
+	URL     string          `json:"url"`
+	Headers http.Header     `json:"headers"`
+	Body    json.RawMessage `json:"body,omitempty"`
+}
+
 type TraceData struct {
-	Version string              `json:"version"`
-	Info    *TraceInfo          `json:"info"`
-	Fetches *FetchTreeTraceNode `json:"fetches"`
+	Version     string              `json:"version"`
+	Info        *TraceInfo          `json:"info"`
+	Fetches     *FetchTreeTraceNode `json:"fetches"`
+	RequestData *RequestData        `json:"request"`
 }
 
 func GetTrace(ctx context.Context, fetchTree *FetchTreeNode) TraceData {
-	return TraceData{
+	trace := TraceData{
 		Version: "1",
 		Info:    GetTraceInfo(ctx),
 		Fetches: fetchTree.Trace(),
 	}
+
+	if req := GetRequest(ctx); req != nil {
+		trace.RequestData = req
+	}
+
+	return trace
 }