Consider abstract lookups

tobias-tengler · tobias-tengler · commit 0e0f3802a34f · 2025-08-02T20:32:21.000+02:00
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceInterfaceTypeCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceInterfaceTypeCollection.cs
@@ -3,16 +3,10 @@
 
 namespace HotChocolate.Fusion.Types.Collections;
 
-public class SourceInterfaceTypeCollection
-    : SourceMemberCollection<SourceInterfaceType>
-    , ISourceComplexTypeCollection<SourceInterfaceType>
-    , ISourceComplexTypeCollection<ISourceComplexType>
+public class SourceInterfaceTypeCollection(IEnumerable<SourceInterfaceType> members)
+    : SourceMemberCollection<SourceInterfaceType>(members), ISourceComplexTypeCollection<SourceInterfaceType>
+        , ISourceComplexTypeCollection<ISourceComplexType>
 {
-    public SourceInterfaceTypeCollection(IEnumerable<SourceInterfaceType> members)
-        : base(members)
-    {
-    }
-
     ISourceComplexType ISourceMemberCollection<ISourceComplexType>.this[string schemaName]
         => this[schemaName];
 
@@ -22,12 +16,9 @@ public ImmutableArray<SourceInterfaceType> Types
     ImmutableArray<ISourceComplexType> ISourceComplexTypeCollection<ISourceComplexType>.Types
         => [.. Members];
 
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out SourceInterfaceType? type)
-        => TryGetMember(schemaName, out type);
-
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
+    public bool TryGetMember(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
     {
-        if (TryGetMember(schemaName, out var member))
+        if (base.TryGetMember(schemaName, out var member))
         {
             type = member;
             return true;
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceMemberCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceMemberCollection.cs
@@ -11,7 +11,7 @@ public class SourceMemberCollection<TMember>
 {
     private readonly FrozenDictionary<string, TMember> _members;
 
-    protected SourceMemberCollection(IEnumerable<TMember> members)
+    public SourceMemberCollection(IEnumerable<TMember> members)
     {
         _members = members.ToFrozenDictionary(t => t.SchemaName);
     }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceObjectTypeCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceObjectTypeCollection.cs
@@ -17,12 +17,9 @@ public ImmutableArray<SourceObjectType> Types
     ImmutableArray<ISourceComplexType> ISourceComplexTypeCollection<ISourceComplexType>.Types
         => [.. Members];
 
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out SourceObjectType? type)
-        => TryGetMember(schemaName, out type);
-
-    public bool TryGetType(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
+    public bool TryGetMember(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)
     {
-        if (TryGetMember(schemaName, out var member))
+        if (base.TryGetMember(schemaName, out var member))
         {
             type = member;
             return true;
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceUnionTypeCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Collections/SourceUnionTypeCollection.cs
@@ -0,0 +1,6 @@
+namespace HotChocolate.Fusion.Types.Collections;
+
+public class SourceUnionTypeCollection(IEnumerable<SourceUnionType> members)
+    : SourceMemberCollection<SourceUnionType>(members)
+{
+}
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompletionTools.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompletionTools.cs
@@ -137,6 +137,29 @@ public static SourceInterfaceTypeCollection CreateSourceInterfaceTypeCollection(
         return new SourceInterfaceTypeCollection(temp);
     }
 
+    public static SourceUnionTypeCollection CreateSourceUnionTypeCollection(
+        UnionTypeDefinitionNode typeDef,
+        CompositeSchemaBuilderContext context)
+    {
+        var types = TypeDirectiveParser.Parse(typeDef.Directives);
+        var lookupDirectives = LookupDirectiveParser.Parse(typeDef.Directives);
+        var temp = new SourceUnionType[types.Length];
+
+        for (var i = 0; i < types.Length; i++)
+        {
+            var type = types[i];
+            var lookups = GetLookupBySchema(lookupDirectives, type.SchemaName, typeDef.Name.Value);
+            context.RegisterForCompletionRange(lookups);
+
+            temp[i] = new SourceUnionType(
+                typeDef.Name.Value,
+                type.SchemaName,
+                lookups);
+        }
+
+        return new SourceUnionTypeCollection(temp);
+    }
+
     private static ImmutableArray<Lookup> GetLookupBySchema(
         ImmutableArray<LookupDirective> allLookups,
         string schemaName,
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeSchemaBuilder.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeSchemaBuilder.cs
@@ -532,7 +532,8 @@ private static void CompleteUnionType(
     {
         var directives = CompletionTools.CreateDirectiveCollection(typeDef.Directives, context);
         var types = CompletionTools.CreateObjectTypeCollection(typeDef.Types, context);
-        type.Complete(new CompositeUnionTypeCompletionContext(types, directives, FeatureCollection.Empty));
+        var sources = CompletionTools.CreateSourceUnionTypeCollection(typeDef, context);
+        type.Complete(new CompositeUnionTypeCompletionContext(types, directives, sources, FeatureCollection.Empty));
     }
 
     private static void CompleteOutputField(
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeUnionTypeCompletionContext.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Completion/CompositeUnionTypeCompletionContext.cs
@@ -6,11 +6,14 @@ namespace HotChocolate.Fusion.Types.Completion;
 internal readonly ref struct CompositeUnionTypeCompletionContext(
     FusionObjectTypeDefinitionCollection types,
     FusionDirectiveCollection directives,
+    SourceUnionTypeCollection sources,
     IFeatureCollection features)
 {
     public FusionDirectiveCollection Directives { get; } = directives;
 
     public FusionObjectTypeDefinitionCollection Types { get; } = types;
 
+    public SourceUnionTypeCollection Sources { get; } = sources;
+
     public IFeatureCollection Features { get; } = features;
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceComplexTypeCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceComplexTypeCollection.cs
@@ -7,7 +7,5 @@ public interface ISourceComplexTypeCollection<TType>
     : ISourceMemberCollection<TType>
     where TType : ISourceComplexType
 {
-    bool TryGetType(string schemaName, [NotNullWhen(true)] out TType? type);
-
     ImmutableArray<TType> Types { get; }
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceMemberCollection.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Contracts/ISourceMemberCollection.cs
@@ -1,8 +1,9 @@
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 
 namespace HotChocolate.Fusion.Types;
 
-public interface ISourceMemberCollection<out TMember>
+public interface ISourceMemberCollection<TMember>
     : IEnumerable<TMember>
     where TMember : ISourceMember
 {
@@ -13,4 +14,6 @@ public interface ISourceMemberCollection<out TMember>
     bool ContainsSchema(string schemaName);
 
     ImmutableArray<string> Schemas { get; }
+
+    bool TryGetMember(string schemaName, [NotNullWhen(true)] out TMember? type);
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/FusionUnionTypeDefinition.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/FusionUnionTypeDefinition.cs
@@ -32,6 +32,23 @@ public FusionUnionTypeDefinition(string name, string? description)
 
     public SchemaCoordinate Coordinate => new(Name, ofDirective: false);
 
+    /// <summary>
+    /// Gets the source type definition of this type.
+    /// </summary>
+    /// <value>
+    /// The source type definition of this type.
+    /// </value>
+    public SourceUnionTypeCollection Sources
+    {
+        get;
+        private set
+        {
+            EnsureNotSealed(_completed);
+
+            field = value;
+        }
+    } = null!;
+
     public FusionObjectTypeDefinitionCollection Types
     {
         get;
@@ -73,6 +90,7 @@ internal void Complete(CompositeUnionTypeCompletionContext context)
 
         Directives = context.Directives;
         Types = context.Types;
+        Sources = context.Sources;
         Features = context.Features;
 
         _completed = true;
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Metadata/SourceUnionType.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/Metadata/SourceUnionType.cs
@@ -0,0 +1,15 @@
+using System.Collections.Immutable;
+
+namespace HotChocolate.Fusion.Types;
+
+public sealed class SourceUnionType(
+    string name,
+    string schemaName,
+    ImmutableArray<Lookup> lookups) : ISourceMember
+{
+    public string Name { get; } = name;
+
+    public string SchemaName { get; } = schemaName;
+
+    public ImmutableArray<Lookup> Lookups { get; } = lookups;
+}
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationPlanner.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Execution/Planning/OperationPlanner.cs
@@ -1039,7 +1039,7 @@ public static void Enqueue(
                 break;
 
             case FieldRequirementWorkItem wi:
-                possiblePlans.EnqueueRequirePlanNodes(planNodeTemplate, wi);
+                possiblePlans.EnqueueRequirePlanNodes(planNodeTemplate, wi, compositeSchema);
                 break;
 
             default:
@@ -1076,7 +1076,7 @@ private static void EnqueueLookupPlanNodes(
 
         foreach (var (schemaName, resolutionCost) in compositeSchema.GetPossibleSchemas(workItem.SelectionSet))
         {
-            foreach (var lookup in workItem.SelectionSet.Type.GetPossibleLookups(schemaName))
+            foreach (var lookup in workItem.SelectionSet.Type.GetPossibleLookups(schemaName, compositeSchema))
             {
                 possiblePlans.Enqueue(
                     planNodeTemplate with
@@ -1093,7 +1093,8 @@ planNodeTemplate with
     private static void EnqueueRequirePlanNodes(
         this PriorityQueue<PlanNode, double> possiblePlans,
         PlanNode planNodeTemplate,
-        FieldRequirementWorkItem workItem)
+        FieldRequirementWorkItem workItem,
+        FusionSchemaDefinition compositeSchema)
     {
         var backlog = planNodeTemplate.Backlog.Pop();
 
@@ -1109,7 +1110,7 @@ planNodeTemplate with
                         Backlog = backlog.Push(workItem)
                     });
 
-                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName))
+                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName, compositeSchema))
                 {
                     possiblePlans.Enqueue(
                         planNodeTemplate with
@@ -1123,7 +1124,7 @@ planNodeTemplate with
             }
             else
             {
-                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName))
+                foreach (var lookup in workItem.Selection.Field.DeclaringType.GetPossibleLookups(schemaName, compositeSchema))
                 {
                     possiblePlans.Enqueue(
                         planNodeTemplate with
@@ -1208,12 +1209,38 @@ static void TrackSchema(Dictionary<string, int> possibleSchemas, string schemaNa
         }
     }
 
-    private static IEnumerable<Lookup> GetPossibleLookups(this ITypeDefinition type, string schemaName)
+    private static IEnumerable<Lookup> GetPossibleLookups(
+        this ITypeDefinition type,
+        string schemaName,
+        FusionSchemaDefinition compositeSchema)
     {
         if (type is FusionComplexTypeDefinition complexType
-            && complexType.Sources.TryGetType(schemaName, out var source))
+            && complexType.Sources.TryGetMember(schemaName, out var source))
         {
-            return source.Lookups;
+            var lookups = source.Lookups.ToList();
+
+            foreach (var interfaceType in complexType.Implements)
+            {
+                if (interfaceType.Sources.TryGetMember(schemaName, out var interfaceSource))
+                {
+                    lookups.AddRange(interfaceSource.Lookups);
+                }
+            }
+
+            var unionTypes = compositeSchema.Types
+                .OfType<FusionUnionTypeDefinition>()
+                .Where(u => u.Types.Contains(type))
+                .ToArray();
+
+            foreach (var unionType in unionTypes)
+            {
+                if (unionType.Sources.TryGetMember(schemaName, out var unionSource))
+                {
+                    lookups.AddRange(unionSource.Lookups);
+                }
+            }
+
+            return lookups;
         }
 
         return [];
diff --git a/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/OperationPlannerTests.cs b/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/OperationPlannerTests.cs
@@ -149,7 +149,7 @@ query GetTopProducts {
     }
 
     [Fact]
-    public void Plan_Simple_Abstract_Lookup()
+    public void Plan_Simple_Interface_Lookup()
     {
         // arrange
         var schema = ComposeSchema(
@@ -159,7 +159,7 @@ schema @schemaName(value: "A") {
             }
 
             type Query {
-              topProducts: [Product!]
+              topProduct: Product
             }
 
             type Product {
@@ -191,7 +191,62 @@ type Product implements Node {
             schema,
             """
             query GetTopProducts {
-              topProducts {
+              topProduct {
+                id
+                name
+                price
+              }
+            }
+            """);
+
+        // assert
+        MatchSnapshot(plan);
+    }
+
+    [Fact]
+    public void Plan_Simple_Union_Lookup()
+    {
+        // arrange
+        var schema = ComposeSchema(
+            """
+            schema @schemaName(value: "A") {
+              query: Query
+            }
+
+            type Query {
+              topProduct: Product
+              # Just here to satisfy satisfiability as I can't make the union lookup internal...
+              productById(id: ID!): Product @lookup @internal
+            }
+
+            type Product {
+              id: ID!
+              name: String!
+            }
+            """,
+            """
+            schema @schemaName(value: "B") {
+              query: Query
+            }
+
+            type Query {
+              lookupUnionById(id: ID!): SomeUnion @lookup
+            }
+
+            union SomeUnion = Product
+
+            type Product {
+              id: ID!
+              price: Float!
+            }
+            """);
+
+        // act
+        var plan = PlanOperation(
+            schema,
+            """
+            query GetTopProducts {
+              topProduct {
                 id
                 name
                 price
diff --git a/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Interface_Lookup.yaml b/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Interface_Lookup.yaml
diff --git a/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Union_Lookup.yaml b/src/HotChocolate/Fusion-vnext/test/Fusion.Execution.Tests/Planning/__snapshots__/OperationPlannerTests.Plan_Simple_Union_Lookup.yaml

Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@ public class SourceMemberCollection<TMember>`
`11`	`11`	`{`
`12`	`12`	`private readonly FrozenDictionary<string, TMember> _members;`
`13`	`13`
`14`		`- protected SourceMemberCollection(IEnumerable<TMember> members)`
	`14`	`+ public SourceMemberCollection(IEnumerable<TMember> members)`
`15`	`15`	`{`
`16`	`16`	`_members = members.ToFrozenDictionary(t => t.SchemaName);`
`17`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -17,12 +17,9 @@ public ImmutableArray<SourceObjectType> Types`
`17`	`17`	`ImmutableArray<ISourceComplexType> ISourceComplexTypeCollection<ISourceComplexType>.Types`
`18`	`18`	`=> [.. Members];`
`19`	`19`
`20`		`- public bool TryGetType(string schemaName, [NotNullWhen(true)] out SourceObjectType? type)`
`21`		`- => TryGetMember(schemaName, out type);`
`22`		`-`
`23`		`- public bool TryGetType(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)`
	`20`	`+ public bool TryGetMember(string schemaName, [NotNullWhen(true)] out ISourceComplexType? type)`
`24`	`21`	`{`
`25`		`- if (TryGetMember(schemaName, out var member))`
	`22`	`+ if (base.TryGetMember(schemaName, out var member))`
`26`	`23`	`{`
`27`	`24`	`type = member;`
`28`	`25`	`return true;`
Original file line number	Diff line number	Diff line change
`@@ -532,7 +532,8 @@ private static void CompleteUnionType(`
`532`	`532`	`{`
`533`	`533`	`var directives = CompletionTools.CreateDirectiveCollection(typeDef.Directives, context);`
`534`	`534`	`var types = CompletionTools.CreateObjectTypeCollection(typeDef.Types, context);`
`535`		`- type.Complete(new CompositeUnionTypeCompletionContext(types, directives, FeatureCollection.Empty));`
	`535`	`+ var sources = CompletionTools.CreateSourceUnionTypeCollection(typeDef, context);`
	`536`	`+ type.Complete(new CompositeUnionTypeCompletionContext(types, directives, sources, FeatureCollection.Empty));`
`536`	`537`	`}`
`537`	`538`
`538`	`539`	`private static void CompleteOutputField(`
Original file line number	Diff line number	Diff line change
`@@ -6,11 +6,14 @@ namespace HotChocolate.Fusion.Types.Completion;`
`6`	`6`	`internal readonly ref struct CompositeUnionTypeCompletionContext(`
`7`	`7`	`FusionObjectTypeDefinitionCollection types,`
`8`	`8`	`FusionDirectiveCollection directives,`
	`9`	`+ SourceUnionTypeCollection sources,`
`9`	`10`	`IFeatureCollection features)`
`10`	`11`	`{`
`11`	`12`	`public FusionDirectiveCollection Directives { get; } = directives;`
`12`	`13`
`13`	`14`	`public FusionObjectTypeDefinitionCollection Types { get; } = types;`
`14`	`15`
	`16`	`+ public SourceUnionTypeCollection Sources { get; } = sources;`
	`17`	`+`
`15`	`18`	`public IFeatureCollection Features { get; } = features;`
`16`	`19`	`}`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,5 @@ public interface ISourceComplexTypeCollection<TType>`
`7`	`7`	`: ISourceMemberCollection<TType>`
`8`	`8`	`where TType : ISourceComplexType`
`9`	`9`	`{`
`10`		`- bool TryGetType(string schemaName, [NotNullWhen(true)] out TType? type);`
`11`		`-`
`12`	`10`	`ImmutableArray<TType> Types { get; }`
`13`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,8 +1,9 @@`
`1`	`1`	`using System.Collections.Immutable;`
	`2`	`+using System.Diagnostics.CodeAnalysis;`
`2`	`3`
`3`	`4`	`namespace HotChocolate.Fusion.Types;`
`4`	`5`
`5`		`-public interface ISourceMemberCollection<out TMember>`
	`6`	`+public interface ISourceMemberCollection<TMember>`
`6`	`7`	`: IEnumerable<TMember>`
`7`	`8`	`where TMember : ISourceMember`
`8`	`9`	`{`
`@@ -13,4 +14,6 @@ public interface ISourceMemberCollection<out TMember>`
`13`	`14`	`bool ContainsSchema(string schemaName);`
`14`	`15`
`15`	`16`	`ImmutableArray<string> Schemas { get; }`
	`17`	`+`
	`18`	`+ bool TryGetMember(string schemaName, [NotNullWhen(true)] out TMember? type);`
`16`	`19`	`}`