[Fusion] Fix BatchExecutionState with multiple VariableValues (#7982)

ChilliCream · Feb 4, 2025 · df2f85d · df2f85d
1 parent cb62bfd
commit df2f85d
Show file tree

Hide file tree

Showing 4 changed files with 577 additions and 3 deletions.
diff --git a/src/HotChocolate/Fusion/src/Core/Execution/Nodes/ResolveByKeyBatch.cs b/src/HotChocolate/Fusion/src/Core/Execution/Nodes/ResolveByKeyBatch.cs
@@ -75,7 +75,7 @@ protected override async Task OnExecuteAsync(
             var batchExecutionState = CreateBatchBatchState(executionState, Requires);
 
             // Create the batch subgraph request.
-            var variableValues = BuildVariables(batchExecutionState, _argumentTypes);
+            var variableValues = BuildVariables(batchExecutionState, Requires, _argumentTypes);
             var request = CreateRequest(context.OperationContext.Variables, variableValues);
 
             // Once we have the batch request, we will enqueue it for execution with
@@ -194,6 +194,7 @@ private void ProcessResult(
 
     private static Dictionary<string, IValueNode> BuildVariables(
         BatchExecutionState[] batchExecutionState,
+        string[] requires,
         Dictionary<string, ITypeNode> argumentTypes)
     {
         var first = batchExecutionState[0];
@@ -220,7 +221,7 @@ private static Dictionary<string, IValueNode> BuildVariables(
             batchState = ref Unsafe.Add(ref batchState, 1)!;
         }
 
-        foreach (var key in first.VariableValues.Keys)
+        foreach (var key in requires)
         {
             var expectedType = argumentTypes[key];
 
@@ -240,7 +241,7 @@ private static Dictionary<string, IValueNode> BuildVariables(
             }
             else
             {
-                if (batchExecutionState[0].VariableValues.TryGetValue(key, out var variableValue))
+                if (first.VariableValues.TryGetValue(key, out var variableValue))
                 {
                     variableValues.Add(key, variableValue);
                 }

diff --git a/src/HotChocolate/Fusion/test/Core.Tests/DemoIntegrationTests.cs b/src/HotChocolate/Fusion/test/Core.Tests/DemoIntegrationTests.cs
@@ -1961,6 +1961,223 @@ query TopProducts {
         Assert.Null(result.ExpectOperationResult().Errors);
     }
 
+    [Fact]
+    public async Task BatchExecutionState_With_Multiple_Variable_Values()
+    {
+        // arrange
+        var subgraphA = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              id: ID!
+              displayName: String!
+            }
+            """);
+        var subgraphB = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+              userBySlug(slug: String!): User
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              relativeUrl: String!
+              id: ID!
+            }
+            """);
+        var subgraphC = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              id: ID!
+              feedbacks: FeedbacksConnection
+            }
+
+            type FeedbacksConnection {
+              edges: [FeedbacksEdge!]
+            }
+
+            type FeedbacksEdge {
+              node: ResaleFeedback!
+            }
+
+            type ResaleFeedback implements Node {
+              feedback: ResaleSurveyFeedback
+              id: ID!
+            }
+
+            type ResaleSurveyFeedback {
+              buyer: User
+            }
+            """);
+
+        using var subgraphs = new TestSubgraphCollection(output, [subgraphA, subgraphB, subgraphC]);
+        var executor = await subgraphs.GetExecutorAsync();
+
+        var request = Parse(
+            """
+            query {
+              userBySlug(slug: "me") {
+                feedbacks {
+                  edges {
+                    node {
+                      feedback {
+                        buyer {
+                          relativeUrl
+                          displayName
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+
+            """);
+
+        // act
+        await using var result = await executor.ExecuteAsync(
+            OperationRequestBuilder
+                .New()
+                .SetDocument(request)
+                .Build());
+
+        // assert
+        var snapshot = new Snapshot();
+        CollectSnapshotData(snapshot, request, result);
+        await snapshot.MatchMarkdownAsync();
+    }
+
+    [Fact]
+    public async Task BatchExecutionState_With_Multiple_Variable_Values_And_Forwarded_Variable()
+    {
+        // arrange
+        var subgraphA = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              id: ID!
+              displayName(arg: String): String!
+            }
+            """);
+        var subgraphB = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+              userBySlug(slug: String!): User
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              relativeUrl(arg: String): String!
+              id: ID!
+            }
+            """);
+        var subgraphC = await TestSubgraph.CreateAsync(
+            """
+            type Query {
+              node(id: ID!): Node
+              nodes(ids: [ID!]!): [Node]!
+            }
+
+            interface Node {
+              id: ID!
+            }
+
+            type User implements Node {
+              id: ID!
+              feedbacks: FeedbacksConnection
+            }
+
+            type FeedbacksConnection {
+              edges: [FeedbacksEdge!]
+            }
+
+            type FeedbacksEdge {
+              node: ResaleFeedback!
+            }
+
+            type ResaleFeedback implements Node {
+              feedback: ResaleSurveyFeedback
+              id: ID!
+            }
+
+            type ResaleSurveyFeedback {
+              buyer: User
+            }
+            """);
+
+        using var subgraphs = new TestSubgraphCollection(output, [subgraphA, subgraphB, subgraphC]);
+        var executor = await subgraphs.GetExecutorAsync();
+
+        var request = Parse(
+            """
+            query($arg1: String, $arg2: String) {
+              userBySlug(slug: "me") {
+                feedbacks {
+                  edges {
+                    node {
+                      feedback {
+                        buyer {
+                          relativeUrl(arg: $arg1)
+                          displayName(arg: $arg2)
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+
+            """);
+
+        // act
+        await using var result = await executor.ExecuteAsync(
+            OperationRequestBuilder
+                .New()
+                .SetDocument(request)
+                .SetVariableValues(new Dictionary<string, object?> { ["arg1"] = "abc", ["arg2"] = "def" })
+                .Build());
+
+        // assert
+        var snapshot = new Snapshot();
+        CollectSnapshotData(snapshot, request, result);
+        await snapshot.MatchMarkdownAsync();
+    }
+
     public sealed class HotReloadConfiguration : IObservable<GatewayConfiguration>
     {
         private GatewayConfiguration _configuration;