feat: implement request tracing (#676)

wundergraph · Nov 20, 2023 · 35355f4 · 35355f4
1 parent 437b781
commit 35355f4
Show file tree

Hide file tree

Showing 20 changed files with 425 additions and 1,476 deletions.
diff --git a/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_federation_test.go b/v2/pkg/engine/datasource/graphql_datasource/graphql_datasource_federation_test.go
@@ -512,6 +512,9 @@ func TestGraphQLDataSourceFederation(t *testing.T) {
 								DataSource:     &Source{},
 								PostProcessing: DefaultPostProcessingConfiguration,
 							},
+							Info: &resolve.FetchInfo{
+								DataSourceID: "user.service",
+							},
 						},
 						Fields: []*resolve.Field{
 							{
@@ -589,7 +592,10 @@ func TestGraphQLDataSourceFederation(t *testing.T) {
 													},
 												},
 												Fetch: &resolve.SingleFetch{
-													SerialID:             1,
+													SerialID: 1,
+													Info: &resolve.FetchInfo{
+														DataSourceID: "account.service",
+													},
 													DataSourceIdentifier: []byte("graphql_datasource.Source"),
 													FetchConfiguration: resolve.FetchConfiguration{
 														Input:                                 `{"method":"POST","url":"http://account.service","body":{"query":"query($representations: [_Any!]!){_entities(representations: $representations){__typename ... on Account {name shippingInfo {zip}}}}","variables":{"representations":[$$0$$]}}}`,

diff --git a/v2/pkg/engine/datasource/httpclient/httpclient.go b/v2/pkg/engine/datasource/httpclient/httpclient.go
@@ -32,6 +32,7 @@ const (
 	UNDEFINED_VARIABLES                         = "undefined"
 	FORWARDED_CLIENT_HEADER_NAMES               = "forwarded_client_header_names"
 	FORWARDED_CLIENT_HEADER_REGULAR_EXPRESSIONS = "forwarded_client_header_regular_expressions"
+	TRACE                                       = "__trace__"
 )
 
 var (
@@ -41,6 +42,7 @@ var (
 		{BODY},
 		{HEADER},
 		{QUERYPARAMS},
+		{TRACE},
 	}
 	subscriptionInputPaths = [][]string{
 		{URL},
@@ -208,7 +210,7 @@ func SetInputPath(input, path []byte) []byte {
 	return out
 }
 
-func requestInputParams(input []byte) (url, method, body, headers, queryParams []byte) {
+func requestInputParams(input []byte) (url, method, body, headers, queryParams []byte, trace bool) {
 	jsonparser.EachKey(input, func(i int, bytes []byte, valueType jsonparser.ValueType, err error) {
 		switch i {
 		case 0:
@@ -221,6 +223,8 @@ func requestInputParams(input []byte) (url, method, body, headers, queryParams [
 			headers = bytes
 		case 4:
 			queryParams = bytes
+		case 5:
+			trace = bytes[0] == 't'
 		}
 	}, inputPaths...)
 	return

diff --git a/v2/pkg/engine/datasource/httpclient/nethttpclient.go b/v2/pkg/engine/datasource/httpclient/nethttpclient.go
@@ -5,6 +5,7 @@ import (
 	"compress/flate"
 	"compress/gzip"
 	"context"
+	"encoding/json"
 	"io"
 	"net/http"
 	"time"
@@ -40,9 +41,26 @@ var (
 	}
 )
 
+type TraceHTTP struct {
+	Request  TraceHTTPRequest  `json:"request"`
+	Response TraceHTTPResponse `json:"response"`
+}
+
+type TraceHTTPRequest struct {
+	Method  string      `json:"method"`
+	URL     string      `json:"url"`
+	Headers http.Header `json:"headers"`
+}
+
+type TraceHTTPResponse struct {
+	StatusCode int         `json:"status_code"`
+	Status     string      `json:"status"`
+	Headers    http.Header `json:"headers"`
+}
+
 func Do(client *http.Client, ctx context.Context, requestInput []byte, out io.Writer) (err error) {
 
-	url, method, body, headers, queryParams := requestInputParams(requestInput)
+	url, method, body, headers, queryParams, enableTrace := requestInputParams(requestInput)
 
 	request, err := http.NewRequestWithContext(ctx, string(method), string(url), bytes.NewReader(body))
 	if err != nil {
@@ -113,8 +131,38 @@ func Do(client *http.Client, ctx context.Context, requestInput []byte, out io.Wr
 		return err
 	}
 
-	_, err = io.Copy(out, respReader)
-	return
+	if !enableTrace {
+		_, err = io.Copy(out, respReader)
+		return
+	}
+
+	buf := &bytes.Buffer{}
+	_, err = io.Copy(buf, respReader)
+	if err != nil {
+		return err
+	}
+	responseTrace := TraceHTTP{
+		Request: TraceHTTPRequest{
+			Method:  request.Method,
+			URL:     request.URL.String(),
+			Headers: request.Header,
+		},
+		Response: TraceHTTPResponse{
+			StatusCode: response.StatusCode,
+			Status:     response.Status,
+			Headers:    response.Header,
+		},
+	}
+	trace, err := json.Marshal(responseTrace)
+	if err != nil {
+		return err
+	}
+	responseWithTraceExtension, err := jsonparser.Set(buf.Bytes(), trace, "extensions", "trace")
+	if err != nil {
+		return err
+	}
+	_, err = out.Write(responseWithTraceExtension)
+	return err
 }
 
 func respBodyReader(resp *http.Response) (io.ReadCloser, error) {

diff --git a/v2/pkg/engine/plan/configuration_visitor.go b/v2/pkg/engine/plan/configuration_visitor.go
@@ -60,6 +60,7 @@ type objectFetchConfiguration struct {
 	fieldRef           int
 	fieldDefinitionRef int
 	fetchID            int
+	sourceID           string
 }
 
 func (c *configurationVisitor) currentSelectionSet() int {
@@ -654,6 +655,7 @@ func (c *configurationVisitor) addNewPlanner(ref int, typeName, fieldName, curre
 		fieldRef:           ref,
 		fieldDefinitionRef: fieldDefinition,
 		fetchID:            fetchID,
+		sourceID:           config.ID,
 	})
 
 	c.saveAddedPath(currentPathConfiguration)

diff --git a/v2/pkg/engine/plan/schemausageinfo.go b/v2/pkg/engine/plan/schemausageinfo.go
@@ -58,15 +58,7 @@ func (t *TypeFieldUsageInfo) Equals(other TypeFieldUsageInfo) bool {
 			return false
 		}
 	}
-	if len(t.Source.IDs) != len(other.Source.IDs) {
-		return false
-	}
-	for i := range t.Source.IDs {
-		if t.Source.IDs[i] != other.Source.IDs[i] {
-			return false
-		}
-	}
-	return true
+	return len(t.Source.IDs) == len(other.Source.IDs)
 }
 
 type InputTypeFieldUsageInfo struct {

diff --git a/v2/pkg/engine/plan/visitor.go b/v2/pkg/engine/plan/visitor.go
@@ -339,6 +339,7 @@ func (v *Visitor) resolveFieldInfo(ref, typeRef int) *resolve.FieldInfo {
 	}
 
 	sourceIDs := make([]string, 0, 1)
+
 	for i := range v.planners {
 		for j := range v.planners[i].paths {
 			if v.planners[i].paths[j].fieldRef == ref {
@@ -1038,5 +1039,11 @@ func (v *Visitor) configureFetch(internal objectFetchConfiguration, external res
 		DataSourceIdentifier: []byte(dataSourceType),
 	}
 
+	if v.Config.IncludeInfo {
+		singleFetch.Info = &resolve.FetchInfo{
+			DataSourceID: internal.sourceID,
+		}
+	}
+
 	return singleFetch
 }
diff --git a/v2/pkg/engine/postprocess/datasourcefetch.go b/v2/pkg/engine/postprocess/datasourcefetch.go
@@ -84,6 +84,7 @@ func (d *DataSourceFetch) createEntityBatchFetch(fetch *resolve.SingleFetch) res
 	}
 
 	return &resolve.BatchEntityFetch{
+		Info: fetch.Info,
 		Input: resolve.BatchInput{
 			Header: resolve.InputTemplate{
 				Segments:                              fetch.InputTemplate.Segments[:representationsVariableIndex],
@@ -128,6 +129,7 @@ func (d *DataSourceFetch) createEntityFetch(fetch *resolve.SingleFetch) resolve.
 	}
 
 	return &resolve.EntityFetch{
+		Info: fetch.Info,
 		Input: resolve.EntityInput{
 			Header: resolve.InputTemplate{
 				Segments:                              fetch.InputTemplate.Segments[:representationsVariableIndex],