Split the apply extensions middleware. (#5151)

ChilliCream · Jun 14, 2022 · e40f6a3 · e40f6a3
1 parent 74f5a82
commit e40f6a3
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 122 deletions.
diff --git a/src/HotChocolate/Stitching/src/Stitching.Types/HotChocolate.Stitching.Types.csproj b/src/HotChocolate/Stitching/src/Stitching.Types/HotChocolate.Stitching.Types.csproj
@@ -16,8 +16,4 @@
     <ProjectReference Include="..\..\..\Core\src\Core\HotChocolate.Core.csproj" />
   </ItemGroup>
 
-  <ItemGroup>
-    <Folder Include="Pipeline\ApplyRemove" />
-  </ItemGroup>
-
 </Project>
diff --git a/...olate/Stitching/src/Stitching.Types/Pipeline/ApplyExtensions/ApplyExtensionsMiddleware.cs b/...olate/Stitching/src/Stitching.Types/Pipeline/ApplyExtensions/ApplyExtensionsMiddleware.cs
@@ -31,18 +31,15 @@ public async ValueTask InvokeAsync(ISchemaMergeContext context)
         var definitions = new Dictionary<string, ITypeSystemDefinitionNode>();
         var extensions = new List<ITypeSystemExtensionNode>();
 
-        foreach (ServiceConfiguration configuration in context.Configurations)
+        for (var i = 0; i < context.Documents.Count; i++)
         {
-            RegisterServiceInfo(extensions, configuration);
+            Document document = context.Documents[i];
+            CollectTypeDefinitions(definitions, extensions, document.SyntaxTree);
+            CollectTypeExtensions(extensions, document.SyntaxTree);
 
-            foreach (DocumentNode document in configuration.Documents)
-            {
-                CollectTypeDefinitions(definitions, extensions, document);
-                CollectTypeExtensions(extensions, document);
-            }
-
-            DocumentNode subgraph = ApplyExtensions(definitions, extensions);
-            context.Documents = context.Documents.Add(new Document(configuration.Name, subgraph));
+            DocumentNode rewritten = ApplyExtensions(definitions, extensions);
+            document = new Document(document.Name, rewritten);
+            context.Documents = context.Documents.SetItem(i, document);
         }
 
         await _next(context);
@@ -88,22 +85,6 @@ public async ValueTask InvokeAsync(ISchemaMergeContext context)
         }
     }
 
-    private static void RegisterServiceInfo(
-        List<ITypeSystemExtensionNode> extensions,
-        ServiceConfiguration configuration)
-    {
-        var serviceName = new DirectiveNode(
-            "_hc_service",
-            new ArgumentNode("name", configuration.Name));
-
-        var serviceInfo = new SchemaExtensionNode(
-            null,
-            new[] { serviceName },
-            Empty<OperationTypeDefinitionNode>());
-
-        extensions.Add(serviceInfo);
-    }
-
     private DocumentNode ApplyExtensions(
         Dictionary<string, ITypeSystemDefinitionNode> definitions,
         List<ITypeSystemExtensionNode> extensions)

diff --git a/...ate/Stitching/src/Stitching.Types/Pipeline/PrepareDocuments/PrepareDocumentsMiddleware.cs b/...ate/Stitching/src/Stitching.Types/Pipeline/PrepareDocuments/PrepareDocumentsMiddleware.cs
@@ -0,0 +1,55 @@
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using HotChocolate.Language;
+
+namespace HotChocolate.Stitching.Types.Pipeline.PrepareDocuments;
+
+public class PrepareDocumentsMiddleware
+{
+    private readonly MergeSchema _next;
+
+    public PrepareDocumentsMiddleware(MergeSchema next)
+    {
+        _next = next;
+    }
+
+    public async ValueTask InvokeAsync(ISchemaMergeContext context)
+    {
+        foreach (ServiceConfiguration configuration in context.Configurations)
+        {
+            var definitions = new List<IDefinitionNode>();
+
+            RegisterServiceInfo(definitions, configuration);
+
+            foreach (DocumentNode document in configuration.Documents)
+            {
+                foreach (IDefinitionNode definition in document.Definitions)
+                {
+                    definitions.Add(definition);
+                }
+            }
+
+            context.Documents = context.Documents.Add(
+                new Document(configuration.Name, new DocumentNode(definitions)));
+        }
+
+        await _next(context);
+    }
+
+    private static void RegisterServiceInfo(
+        List<IDefinitionNode> definitions,
+        ServiceConfiguration configuration)
+    {
+        var serviceName = new DirectiveNode(
+            "_hc_service",
+            new ArgumentNode("name", configuration.Name));
+
+        var serviceInfo = new SchemaExtensionNode(
+            null,
+            new[] { serviceName },
+            Array.Empty<OperationTypeDefinitionNode>());
+
+        definitions.Add(serviceInfo);
+    }
+}
diff --git a/...Chocolate/Stitching/test/Stitching.Types.Tests/Pipeline/ApplyExtensionsMiddlewareTests.cs b/...Chocolate/Stitching/test/Stitching.Types.Tests/Pipeline/ApplyExtensionsMiddlewareTests.cs
@@ -2,6 +2,8 @@
 using System.Linq;
 using System.Threading.Tasks;
 using HotChocolate.Stitching.Types.Pipeline.ApplyExtensions;
+using HotChocolate.Stitching.Types.Pipeline.ApplyRenaming;
+using HotChocolate.Stitching.Types.Pipeline.PrepareDocuments;
 using Snapshooter.Xunit;
 using Xunit;
 using static HotChocolate.Language.Utf8GraphQLParser;
@@ -14,7 +16,7 @@ public class ApplyExtensionsMiddlewareTests
     public async Task Apply_Object_Extension_Single_Document()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -30,7 +32,7 @@ public async Task Apply_Object_Extension_Single_Document()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
@@ -40,7 +42,7 @@ public async Task Apply_Object_Extension_Single_Document()
     public async Task Apply_Object_Extension_Is_Preserved()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -56,7 +58,7 @@ public async Task Apply_Object_Extension_Is_Preserved()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
@@ -66,7 +68,7 @@ public async Task Apply_Object_Extension_Is_Preserved()
     public async Task Apply_Object_Extension_Merge_Field_Directives_Single_Document()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -82,7 +84,7 @@ public async Task Apply_Object_Extension_Merge_Field_Directives_Single_Document(
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
@@ -92,7 +94,7 @@ public async Task Apply_Object_Extension_Merge_Field_Directives_Single_Document(
     public async Task Apply_Object_Extension_Merge_Directives_Single_Document()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -106,7 +108,7 @@ public async Task Apply_Object_Extension_Merge_Directives_Single_Document()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
@@ -116,7 +118,7 @@ public async Task Apply_Object_Extension_Merge_Directives_Single_Document()
     public async Task Apply_Object_Extension_Merge_Directives_2_Single_Document()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -130,7 +132,7 @@ public async Task Apply_Object_Extension_Merge_Directives_2_Single_Document()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
@@ -140,7 +142,7 @@ public async Task Apply_Object_Extension_Merge_Directives_2_Single_Document()
     public async Task Apply_Object_Extension_Field_Type_Mismatch()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -156,7 +158,7 @@ public async Task Apply_Object_Extension_Field_Type_Mismatch()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        async Task Error() => await middleware.InvokeAsync(context);
+        async Task Error() => await pipeline(context);
 
         // assert
         await Assert.ThrowsAsync<GraphQLException>(Error);
@@ -166,7 +168,7 @@ public async Task Apply_Object_Extension_Field_Type_Mismatch()
     public async Task Apply_Local_Remove()
     {
         // arrange
-        var middleware = new ApplyExtensionsMiddleware(_ => default);
+        MergeSchema pipeline = CreatePipeline();
 
         var service = new ServiceConfiguration(
             "abc",
@@ -182,9 +184,23 @@ public async Task Apply_Local_Remove()
         var context = new SchemaMergeContext(configurations);
 
         // act
-        await middleware.InvokeAsync(context);
+        await pipeline(context);
 
         // assert
         context.Documents.Single().SyntaxTree.ToString().MatchSnapshot();
     }
+
+    private MergeSchema CreatePipeline()
+        => new SchemaMergePipelineBuilder()
+            .Use(next =>
+            {
+                var middleware = new PrepareDocumentsMiddleware(next);
+                return context => middleware.InvokeAsync(context);
+            })
+            .Use(next =>
+            {
+                var middleware = new ApplyExtensionsMiddleware(next);
+                return context => middleware.InvokeAsync(context);
+            })
+            .Compile();
 }