fix: preserve existing context metadata

Author: dkorittki

v2/pkg/engine/datasource/grpc_datasource/grpc_datasource.go

@@ -117,11 +117,15 @@ func (d *DataSource) Load(ctx context.Context, headers http.Header, input []byte
 
 	// convert headers to grpc metadata and attach to ctx
 	if len(headers) > 0 {
-		md := make(metadata.MD, len(headers))
-		for k, v := range headers {
-			md.Set(strings.ToLower(k), v...)
+		// assume that each header has exactly one value for default pairs size
+		pairs := make([]string, 0, len(headers)*2)
+		for headerName, headerValues := range headers {
+			headerName = strings.ToLower(headerName)
+			for _, v := range headerValues {
+				pairs = append(pairs, headerName, v)
+			}
 		}
-		ctx = metadata.NewOutgoingContext(ctx, md)
+		ctx = metadata.AppendToOutgoingContext(ctx, pairs...)
 	}
 
 	graph := NewDependencyGraph(d.plan)

v2/pkg/engine/datasource/grpc_datasource/grpc_datasource_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/tidwall/gjson"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/test/bufconn"
 	"google.golang.org/protobuf/encoding/protojson"
 	protoref "google.golang.org/protobuf/reflect/protoreflect"
@@ -5445,3 +5446,69 @@ func Test_Datasource_Load_WithHeaders(t *testing.T) {
 		})
 	}
 }
+
+func Test_Datasource_Load_PreservesExistingContextMetadata(t *testing.T) {
+	conn, cleanup := setupTestGRPCServer(t)
+	t.Cleanup(cleanup)
+
+	type graphqlError struct {
+		Message string `json:"message"`
+	}
+	type graphqlResponse struct {
+		Data   map[string]interface{} `json:"data"`
+		Errors []graphqlError         `json:"errors,omitempty"`
+	}
+
+	// Parse the GraphQL schema
+	schemaDoc := grpctest.MustGraphQLSchema(t)
+
+	query := `query UserQuery($id: ID!) { user(id: $id) { id name } }`
+	vars := `{"variables":{"id":"test-user-123"}}`
+
+	// Parse the GraphQL query
+	queryDoc, report := astparser.ParseGraphqlDocumentString(query)
+	require.False(t, report.HasErrors(), "failed to parse query: %s", report.Error())
+
+	compiler, err := NewProtoCompiler(grpctest.MustProtoSchema(t), testMapping())
+	require.NoError(t, err)
+
+	// Create the datasource
+	ds, err := NewDataSource(conn, DataSourceConfig{
+		Operation:    &queryDoc,
+		Definition:   &schemaDoc,
+		SubgraphName: "Products",
+		Mapping:      testMapping(),
+		Compiler:     compiler,
+	})
+	require.NoError(t, err)
+
+	// Create a context with existing metadata
+	ctx := metadata.NewOutgoingContext(
+		context.Background(),
+		metadata.Pairs("x-existing-key", "existing-value"),
+	)
+
+	// Create HTTP headers to be forwarded
+	headers := make(http.Header)
+	headers.Set("X-User-ID", "header-user-456")
+
+	// Execute the query with both existing context metadata and new HTTP headers
+	input := fmt.Sprintf(`{"query":%q,"body":%s}`, query, vars)
+	output, err := ds.Load(ctx, headers, []byte(input))
+	require.NoError(t, err)
+
+	// Parse the response
+	var resp graphqlResponse
+	err = json.Unmarshal(output, &resp)
+	require.NoError(t, err, "Failed to unmarshal response")
+
+	// Verify no errors
+	require.Empty(t, resp.Errors, "Should not have GraphQL errors")
+
+	// Verify the response includes both the header-derived ID and the existing metadata value
+	user, ok := resp.Data["user"].(map[string]interface{})
+	require.True(t, ok, "user should be an object")
+	require.Equal(t, "header-user-456", user["id"], "user ID should come from HTTP header")
+	require.Equal(t, "User header-user-456 (existing: existing-value)", user["name"],
+		"user name should include both header-derived ID and existing context metadata")
+}

v2/pkg/grpctest/mockservice.go

@@ -236,22 +236,33 @@ func (s *MockService) QueryUsers(ctx context.Context, in *productv1.QueryUsersRe
 
 func (s *MockService) QueryUser(ctx context.Context, in *productv1.QueryUserRequest) (*productv1.QueryUserResponse, error) {
 	userId := in.GetId()
+	existingValue := ""
 	md, ok := metadata.FromIncomingContext(ctx)
 	if ok {
 		if values := md.Get("x-user-id"); len(values) > 0 {
 			userId = values[0]
 		}
+		// Check for existing metadata that was in the context before headers were added
+		if values := md.Get("x-existing-key"); len(values) > 0 {
+			existingValue = values[0]
+		}
 	}
 
 	// Return a gRPC status error for a specific test case
 	if userId == "error-user" {
 		return nil, status.Errorf(codes.NotFound, "user not found: %s", userId)
 	}
 
+	// Include existing metadata value in the name if present
+	userName := fmt.Sprintf("User %s", userId)
+	if existingValue != "" {
+		userName = fmt.Sprintf("User %s (existing: %s)", userId, existingValue)
+	}
+
 	return &productv1.QueryUserResponse{
 		User: &productv1.User{
 			Id:   userId,
-			Name: fmt.Sprintf("User %s", userId),
+			Name: userName,
 		},
 	}, nil
 }