Update selection path with type refinement

Author: tobias-tengler

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationDefinitionBuilder.cs

@@ -12,6 +12,7 @@ internal sealed class OperationDefinitionBuilder
     private string? _name;
     private string? _description;
     private Lookup? _lookup;
+    private string? _lookupTypeRefinement;
     private string? _requirementKey;
     private SelectionSetNode? _selectionSet;
 
@@ -40,9 +41,10 @@ public OperationDefinitionBuilder SetDescription(string? description)
         return this;
     }
 
-    public OperationDefinitionBuilder SetLookup(Lookup? lookup, string? requirementKey)
+    public OperationDefinitionBuilder SetLookup(Lookup? lookup, string? lookupTypeRefinement, string? requirementKey)
     {
         _lookup = lookup;
+        _lookupTypeRefinement = lookupTypeRefinement;
         _requirementKey = requirementKey;
         return this;
     }
@@ -87,6 +89,12 @@ public OperationDefinitionBuilder SetSelectionSet(SelectionSetNode selectionSet)
             var indexBuilder = index.ToBuilder();
             indexBuilder.Register(selectionSet);
             index = indexBuilder;
+
+            if (!string.IsNullOrEmpty(_lookupTypeRefinement))
+            {
+                selectionPath = selectionPath.AppendFragment(_lookupTypeRefinement);
+            }
+
             selectionPath = selectionPath.AppendField(_lookup.FieldName);
         }
 

src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationPlanner.cs

@@ -293,10 +293,11 @@ lookup is null
                 requirements = requirements.Add(argumentRequirementKey, operationRequirement);
             }
 
-            // If the lookup returns an abstract type, we might need to insert a type refinement
-            // around our selection set.
+            string? lookupTypeRefinement = null;
             if (_schema.Types.TryGetType(lookup.FieldType, out var lookupFieldType)
-                && lookupFieldType != workItem.SelectionSet.Type)
+                && lookupFieldType != workItem.SelectionSet.Type
+                && !resolvable.Selections.All(s => s is InlineFragmentNode inlineFragment
+                    && inlineFragment.TypeCondition?.Name.Value == workItem.SelectionSet.Type.Name))
             {
                 var typeRefinement =
                     new InlineFragmentNode(
@@ -313,9 +314,11 @@ lookup is null
                 index = indexBuilder;
 
                 operationBuilder.SetSelectionSet(selectionSetWithTypeRefinement);
+
+                lookupTypeRefinement = workItem.SelectionSet.Type.Name;
             }
 
-            operationBuilder.SetLookup(lookup, requirementKey);
+            operationBuilder.SetLookup(lookup, lookupTypeRefinement, requirementKey);
         }
 
         (var definition, index, var source) = operationBuilder.Build(index);
@@ -686,7 +689,32 @@ private void PlanFieldWithRequirement(
             requirements = requirements.Add(argumentRequirementKey, operationRequirement);
         }
 
-        operationBuilder.SetLookup(workItem.Lookup, requirementKey);
+        string? lookupTypeRefinement = null;
+        if (_schema.Types.TryGetType(lookup.FieldType, out var lookupFieldType)
+            && lookupFieldType != selectionSetStub.Type
+            && !selectionSetNode.Selections.All(s => s is InlineFragmentNode inlineFragment
+                && inlineFragment.TypeCondition?.Name.Value == selectionSetStub.Type.Name))
+        {
+            var typeRefinement =
+                new InlineFragmentNode(
+                    null,
+                    new NamedTypeNode(selectionSetStub.Type.Name),
+                    [],
+                    selectionSetNode);
+            var selectionSetWithTypeRefinement = new SelectionSetNode(null, [typeRefinement]);
+
+            var indexBuilder = index.ToBuilder();
+
+            indexBuilder.Register(selectionSetNode, selectionSetWithTypeRefinement);
+
+            index = indexBuilder;
+
+            operationBuilder.SetSelectionSet(selectionSetWithTypeRefinement);
+
+            lookupTypeRefinement = selectionSetStub.Type.Name;
+        }
+
+        operationBuilder.SetLookup(workItem.Lookup, lookupTypeRefinement, requirementKey);
 
         var (definition, _, source) = operationBuilder.Build(index);
 

src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Interface_Lookup.yaml

@@ -1,13 +1,13 @@
 operation:
   - document: >-
-    query GetTopProducts {
-      topProduct {
-        id
-        name
-        price
-        id @fusion__requirement
+      query GetTopProducts {
+        topProduct {
+          id
+          name
+          price
+          id @fusion__requirement
+        }
       }
-    }
     name: GetTopProducts
     hash: 123
 nodes:
@@ -32,7 +32,7 @@ nodes:
           }
         }
       }
-    source: $.node
+    source: $.node.<Product>
     target: $.topProduct
     requirements:
       - name: __fusion_1_id

src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Union_Lookup.yaml

@@ -1,13 +1,13 @@
 operation:
   - document: >-
-    query GetTopProducts {
-      topProduct {
-        id
-        name
-        price
-        id @fusion__requirement
+      query GetTopProducts {
+        topProduct {
+          id
+          name
+          price
+          id @fusion__requirement
+        }
       }
-    }
     name: GetTopProducts
     hash: 123
 nodes:
@@ -32,7 +32,7 @@ nodes:
           }
         }
       }
-    source: $.lookupUnionById
+    source: $.lookupUnionById.<Product>
     target: $.topProduct
     requirements:
       - name: __fusion_1_id