[Fusion] Applied inferred key directives (#8717)

glen-84 · web-flow · commit 1529d938fa37 · 2025-09-17T09:56:41.000+02:00
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/EnumerableCompositionResultExtensions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/EnumerableCompositionResultExtensions.cs
@@ -0,0 +1,19 @@
+using System.Collections.Immutable;
+using HotChocolate.Fusion.Results;
+
+namespace HotChocolate.Fusion.Extensions;
+
+internal static class EnumerableCompositionResultExtensions
+{
+    public static CompositionResult Combine(this IEnumerable<CompositionResult> results)
+    {
+        var failedResults = results.Where(r => r.IsFailure).ToImmutableArray();
+
+        if (failedResults.Length == 0)
+        {
+            return CompositionResult.Success();
+        }
+
+        return failedResults.SelectMany(r => r.Errors).ToImmutableArray();
+    }
+}
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableComplexTypeDefinitionExtensions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableComplexTypeDefinitionExtensions.cs
@@ -1,17 +1,27 @@
+using HotChocolate.Language;
 using HotChocolate.Types;
 using HotChocolate.Types.Mutable;
 using ArgumentNames = HotChocolate.Fusion.WellKnownArgumentNames;
+using DirectiveNames = HotChocolate.Fusion.WellKnownDirectiveNames;
 
 namespace HotChocolate.Fusion.Extensions;
 
 internal static class MutableComplexTypeDefinitionExtensions
 {
-    public static void ApplyKeyDirective(this MutableComplexTypeDefinition type, string[] fields)
+    public static void ApplyKeyDirective(this MutableComplexTypeDefinition type, string keyFields)
     {
-        var fieldsArgument = new ArgumentAssignment(ArgumentNames.Fields, string.Join(" ", fields));
-        var keyDirective =
-            new Directive(new MutableDirectiveDefinition(WellKnownDirectiveNames.Key), fieldsArgument);
+        var keyDirectiveExists =
+            type.Directives.AsEnumerable().Any(
+                d =>
+                    d.Name == DirectiveNames.Key
+                    && ((StringValueNode)d.Arguments[ArgumentNames.Fields]).Value.Equals(keyFields));
 
-        type.Directives.Add(keyDirective);
+        if (!keyDirectiveExists)
+        {
+            type.Directives.Add(
+                new Directive(
+                    FusionBuiltIns.SourceSchemaDirectives[DirectiveNames.Key],
+                    new ArgumentAssignment(ArgumentNames.Fields, keyFields)));
+        }
     }
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableObjectTypeDefinitionExtensions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableObjectTypeDefinitionExtensions.cs
@@ -5,11 +5,6 @@ namespace HotChocolate.Fusion.Extensions;
 
 internal static class MutableObjectTypeDefinitionExtensions
 {
-    public static void ApplyShareableDirective(this MutableObjectTypeDefinition type)
-    {
-        type.Directives.Add(new Directive(new MutableDirectiveDefinition(DirectiveNames.Shareable)));
-    }
-
     public static bool HasInternalDirective(this MutableObjectTypeDefinition type)
     {
         return type.Directives.ContainsName(DirectiveNames.Internal);
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableOutputFieldDefinitionExtensions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableOutputFieldDefinitionExtensions.cs
@@ -1,4 +1,9 @@
+using System.Collections.Immutable;
+using HotChocolate.Fusion.Definitions;
+using HotChocolate.Fusion.Language;
+using HotChocolate.Fusion.Rewriters;
 using HotChocolate.Language;
+using HotChocolate.Types;
 using HotChocolate.Types.Mutable;
 using static HotChocolate.Language.Utf8GraphQLParser.Syntax;
 using ArgumentNames = HotChocolate.Fusion.WellKnownArgumentNames;
@@ -8,9 +13,14 @@ namespace HotChocolate.Fusion.Extensions;
 
 internal static class MutableOutputFieldDefinitionExtensions
 {
-    public static void ApplyLookupDirective(this MutableOutputFieldDefinition field)
+    public static void ApplyShareableDirective(this MutableOutputFieldDefinition field)
     {
-        field.Directives.Add(new Directive(new MutableDirectiveDefinition(DirectiveNames.Lookup)));
+        if (field.Directives.ContainsName(DirectiveNames.Shareable))
+        {
+            return;
+        }
+
+        field.Directives.Add(new Directive(new ShareableMutableDirectiveDefinition()));
     }
 
     public static bool ExistsInSchema(this MutableOutputFieldDefinition field, string schemaName)
@@ -39,6 +49,23 @@ public static bool ExistsInSchema(this MutableOutputFieldDefinition field, strin
         return null;
     }
 
+    // productById(id: ID!) -> ["id"].
+    // productByIdAndCategoryId(id: ID!, categoryId: Int) -> ["id", "categoryId"].
+    // personByAddressId(id: ID! @is(field: "address.id")) -> ["address.id"].
+    public static List<string> GetFusionLookupMap(this MutableOutputFieldDefinition field)
+    {
+        var items = new List<string>();
+
+        foreach (var argument in field.Arguments)
+        {
+            var @is = argument.GetIsFieldSelectionMap();
+
+            items.Add(@is ?? argument.Name);
+        }
+
+        return items;
+    }
+
     public static SelectionSetNode? GetFusionRequiresRequirements(
         this MutableOutputFieldDefinition field,
         string schemaName)
@@ -61,6 +88,24 @@ public static bool ExistsInSchema(this MutableOutputFieldDefinition field, strin
         return ParseSelectionSet($"{{ {requirements} }}");
     }
 
+    public static string GetKeyFields(
+        this MutableOutputFieldDefinition field,
+        List<string> lookupMap,
+        MutableSchemaDefinition schema)
+    {
+        var selectedValues = lookupMap.Select(a => new FieldSelectionMapParser(a).Parse());
+        var valueSelectionToSelectionSetRewriter = new ValueSelectionToSelectionSetRewriter(schema);
+        var fieldType = field.Type.AsTypeDefinition();
+        var selectionSets = selectedValues
+            .Select(s => valueSelectionToSelectionSetRewriter.Rewrite(s, fieldType))
+            .ToImmutableArray();
+        var mergedSelectionSet = selectionSets.Length == 1
+            ? selectionSets[0]
+            : new MergeSelectionSetRewriter(schema).Merge(selectionSets, fieldType);
+
+        return mergedSelectionSet.ToString(indented: false).AsSpan()[2..^2].ToString();
+    }
+
     public static bool HasInternalDirective(this MutableOutputFieldDefinition type)
     {
         return type.Directives.ContainsName(DirectiveNames.Internal);
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableSchemaDefinitionExtensions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Extensions/MutableSchemaDefinitionExtensions.cs
@@ -9,15 +9,15 @@ internal static class MutableSchemaDefinitionExtensions
 {
     public static void AddBuiltInFusionTypes(this MutableSchemaDefinition schema)
     {
-        foreach (var builtInScalar in FusionBuiltIns.SourceSchemaScalars)
+        foreach (var builtInScalar in FusionBuiltIns.SourceSchemaScalars.Values)
         {
             schema.Types.Add(builtInScalar);
         }
     }
 
     public static void AddBuiltInFusionDirectives(this MutableSchemaDefinition schema)
     {
-        foreach (var builtInDirective in FusionBuiltIns.SourceSchemaDirectives)
+        foreach (var builtInDirective in FusionBuiltIns.SourceSchemaDirectives.Values)
         {
             schema.DirectiveDefinitions.Add(builtInDirective);
         }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/FusionBuiltIns.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/FusionBuiltIns.cs
@@ -15,7 +15,7 @@ internal static class FusionBuiltIns
 
     private static readonly MutableScalarTypeDefinition s_stringType = BuiltIns.String.Create();
 
-    public static FrozenSet<MutableDirectiveDefinition> SourceSchemaDirectives { get; } =
+    public static FrozenDictionary<string, MutableDirectiveDefinition> SourceSchemaDirectives { get; } =
         new HashSet<MutableDirectiveDefinition>(
         [
             new ExternalMutableDirectiveDefinition(),
@@ -28,11 +28,11 @@ internal static class FusionBuiltIns
             new ProvidesMutableDirectiveDefinition(s_fieldSelectionSetType),
             new RequireMutableDirectiveDefinition(s_fieldSelectionMapType),
             new ShareableMutableDirectiveDefinition()
-        ]).ToFrozenSet();
+        ]).ToFrozenDictionary(d => d.Name);
 
-    public static FrozenSet<MutableScalarTypeDefinition> SourceSchemaScalars { get; } =
+    public static FrozenDictionary<string, MutableScalarTypeDefinition> SourceSchemaScalars { get; } =
         new HashSet<MutableScalarTypeDefinition>([s_fieldSelectionMapType, s_fieldSelectionSetType])
-            .ToFrozenSet();
+            .ToFrozenDictionary(s => s.Name);
 
     public static bool IsBuiltInSourceSchemaScalar(string typeName)
     {
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Options/SchemaComposerOptions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Options/SchemaComposerOptions.cs
@@ -3,4 +3,6 @@ namespace HotChocolate.Fusion.Options;
 public sealed class SchemaComposerOptions
 {
     public required bool EnableGlobalObjectIdentification { get; init; }
+
+    public SourceSchemaPreprocessorOptions Preprocessor { get; init; } = new();
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Options/SourceSchemaPreprocessorOptions.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/Options/SourceSchemaPreprocessorOptions.cs
@@ -1,6 +1,9 @@
 namespace HotChocolate.Fusion.Options;
 
-internal sealed class SourceSchemaPreprocessorOptions
+public sealed class SourceSchemaPreprocessorOptions
 {
-    public bool ApplyShareableToAllTypes { get; set; } = true;
+    /// <summary>
+    /// Applies inferred key directives to types that are returned by lookup fields.
+    /// </summary>
+    public bool ApplyInferredKeyDirectives { get; set; } = true;
 }
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SchemaComposer.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SchemaComposer.cs
@@ -1,4 +1,5 @@
 using System.Collections.Immutable;
+using HotChocolate.Fusion.Extensions;
 using HotChocolate.Fusion.Logging.Contracts;
 using HotChocolate.Fusion.Options;
 using HotChocolate.Fusion.PostMergeValidationRules;
@@ -40,6 +41,16 @@ public CompositionResult<MutableSchemaDefinition> Compose()
             return parseErrors;
         }
 
+        // Preprocess Source Schemas
+        var preprocessorOptions = _schemaComposerOptions.Preprocessor;
+        var preprocessResult =
+            schemas.Select(schema => new SourceSchemaPreprocessor(schema, preprocessorOptions).Process()).Combine();
+
+        if (preprocessResult.IsFailure)
+        {
+            return preprocessResult;
+        }
+
         // Validate Source Schemas
         var validationResult =
             new SourceSchemaValidator(schemas, s_sourceSchemaRules, _log).Validate();
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SourceSchemaMerger.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SourceSchemaMerger.cs
@@ -1011,19 +1011,8 @@ private void AddFusionLookupDirectives(
         foreach (var (sourceField, sourcePath, sourceSchema) in lookupFieldGroup)
         {
             var schemaArgument = new EnumValueNode(_schemaConstantNames[sourceSchema]);
-            var lookupMap = GetFusionLookupMap(sourceField);
-            var selectedValues = lookupMap.Select(a => new FieldSelectionMapParser(a).Parse());
-            var selectedValueToSelectionSetRewriter =
-                GetSelectedValueToSelectionSetRewriter(sourceSchema);
-            var selectionSets = selectedValues
-                .Select(
-                    s => selectedValueToSelectionSetRewriter.Rewrite(s, type))
-                .ToImmutableArray();
-            var mergedSelectionSet = selectionSets.Length == 1
-                ? selectionSets[0]
-                : GetMergeSelectionSetRewriter(sourceSchema).Merge(selectionSets, type);
-            var keyArgument =
-                mergedSelectionSet.ToString(indented: false).AsSpan()[2..^2].ToString();
+            var lookupMap = sourceField.GetFusionLookupMap();
+            var keyArgument = sourceField.GetKeyFields(lookupMap, sourceSchema);
 
             var fieldArgument =
                 s_fieldDefinitionRewriter
@@ -1050,23 +1039,6 @@ private void AddFusionLookupDirectives(
         }
     }
 
-    // productById(id: ID!) -> ["id"].
-    // productByIdAndCategoryId(id: ID!, categoryId: Int) -> ["id", "categoryId"].
-    // personByAddressId(id: ID! @is(field: "address.id")) -> ["address.id"].
-    private static List<string> GetFusionLookupMap(MutableOutputFieldDefinition field)
-    {
-        var items = new List<string>();
-
-        foreach (var argument in field.Arguments)
-        {
-            var @is = argument.GetIsFieldSelectionMap();
-
-            items.Add(@is ?? argument.Name);
-        }
-
-        return items;
-    }
-
     private void AddFusionRequiresDirectives(
         MutableOutputFieldDefinition field,
         MutableComplexTypeDefinition complexType,
diff --git a/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SourceSchemaPreprocessor.cs b/src/HotChocolate/Fusion-vnext/src/Fusion.Composition/SourceSchemaPreprocessor.cs
@@ -1,102 +1,57 @@
+using System.Collections.Immutable;
 using HotChocolate.Fusion.Extensions;
 using HotChocolate.Fusion.Options;
+using HotChocolate.Fusion.Results;
 using HotChocolate.Types;
 using HotChocolate.Types.Mutable;
 
 namespace HotChocolate.Fusion;
 
 /// <summary>
-/// Applies @lookup, @key, and optionally @shareable to a source schema to make it equivalent to a Fusion v1 source schema.
+/// Preprocesses a source schema.
 /// </summary>
 internal sealed class SourceSchemaPreprocessor(
     MutableSchemaDefinition schema,
     SourceSchemaPreprocessorOptions? options = null)
 {
-    public MutableSchemaDefinition Process()
-    {
-        var context = new SourceSchemaPreprocessorContext(schema, options ?? new());
-
-        ApplyLookups(context);
+    private readonly SourceSchemaPreprocessorOptions _options = options ?? new SourceSchemaPreprocessorOptions();
 
-        if (context.Options.ApplyShareableToAllTypes)
+    public CompositionResult Process()
+    {
+        if (_options.ApplyInferredKeyDirectives)
         {
-            ApplyShareableToAllTypes(context);
+            ApplyInferredKeyDirectives();
         }
 
-        return schema;
+        return CompositionResult.Success();
     }
 
-    private static void ApplyLookups(SourceSchemaPreprocessorContext context)
+    /// <summary>
+    /// Applies inferred key directives to types that are returned by lookup fields.
+    /// </summary>
+    private void ApplyInferredKeyDirectives()
     {
-        if (context.Schema.QueryType is not { } queryType)
-        {
-            return;
-        }
+        var lookupFieldDefinitions =
+            schema.Types
+                .OfType<MutableComplexTypeDefinition>()
+                .SelectMany(t => t.Fields.AsEnumerable().Where(f => f.HasLookupDirective()));
 
-        foreach (var queryField in queryType.Fields)
+        foreach (var lookupFieldDefinition in lookupFieldDefinitions)
         {
-            if (queryField.HasLookupDirective()
-                || queryField.Type.IsListType()
-                || queryField.Type.Kind == TypeKind.NonNull
-                || queryField.Arguments.Count != 1)
-            {
-                continue;
-            }
-
-            var keyArgument = queryField.Arguments.AsEnumerable().First();
-            var @is = keyArgument.GetIsFieldSelectionMap();
-
-            var queryFieldType = queryField.Type.AsTypeDefinition();
-            var keyOutputFieldName = @is ?? keyArgument.Name;
+            var fieldType = lookupFieldDefinition.Type.AsTypeDefinition();
+            var possibleTypes = schema.GetPossibleTypes(fieldType);
+            var lookupMap = lookupFieldDefinition.GetFusionLookupMap();
+            var keyFields = lookupFieldDefinition.GetKeyFields(lookupMap, schema);
 
-            if (!context.Schema.Types.TryGetType(queryFieldType.Name, out var resultType))
+            foreach (var possibleType in possibleTypes)
             {
-                continue;
+                possibleType.ApplyKeyDirective(keyFields);
             }
 
-            if (resultType is MutableObjectTypeDefinition resultObjectType)
-            {
-                if (resultObjectType.Fields.TryGetField(keyOutputFieldName, out var keyField))
-                {
-                    queryField.ApplyLookupDirective();
-
-                    resultObjectType.ApplyKeyDirective([keyField.Name]);
-                }
-            }
-            else if (resultType is MutableInterfaceTypeDefinition resultInterfaceType)
+            if (fieldType is MutableInterfaceTypeDefinition interfaceType)
             {
-                if (resultInterfaceType.Fields.TryGetField(keyOutputFieldName, out var keyField))
-                {
-                    queryField.ApplyLookupDirective();
-
-                    resultInterfaceType.ApplyKeyDirective([keyField.Name]);
-
-                    foreach (var objectType in context.Schema.Types.OfType<MutableObjectTypeDefinition>())
-                    {
-                        if (objectType.Implements.ContainsName(resultInterfaceType.Name))
-                        {
-                            objectType.ApplyKeyDirective([keyField.Name]);
-                        }
-                    }
-                }
+                interfaceType.ApplyKeyDirective(keyFields);
             }
         }
     }
-
-    private static void ApplyShareableToAllTypes(SourceSchemaPreprocessorContext context)
-    {
-        foreach (var objectType in context.Schema.Types.OfType<MutableObjectTypeDefinition>())
-        {
-            objectType.ApplyShareableDirective();
-        }
-    }
-}
-
-internal sealed class SourceSchemaPreprocessorContext(
-    MutableSchemaDefinition schema,
-    SourceSchemaPreprocessorOptions options)
-{
-    public MutableSchemaDefinition Schema => schema;
-
-    public SourceSchemaPreprocessorOptions Options => options;
 }
diff --git a/src/HotChocolate/Fusion-vnext/test/Fusion.Composition.Tests/SourceSchemaPreprocessorTests.cs b/src/HotChocolate/Fusion-vnext/test/Fusion.Composition.Tests/SourceSchemaPreprocessorTests.cs
diff --git a/src/HotChocolate/Fusion-vnext/test/Fusion.Composition.Tests/__snapshots__/SourceSchemaPreprocessorTests.Preprocess_ApplyInferredKeyDirectivesEnabled_AppliesInferredKeyDirectives.graphql b/src/HotChocolate/Fusion-vnext/test/Fusion.Composition.Tests/__snapshots__/SourceSchemaPreprocessorTests.Preprocess_ApplyInferredKeyDirectivesEnabled_AppliesInferredKeyDirectives.graphql
diff --git a/src/HotChocolate/Fusion-vnext/test/Xunit2.Directory.Build.props b/src/HotChocolate/Fusion-vnext/test/Xunit2.Directory.Build.props
diff --git a/src/HotChocolate/Fusion-vnext/test/Xunit3.Directory.Build.props b/src/HotChocolate/Fusion-vnext/test/Xunit3.Directory.Build.props
diff --git a/src/HotChocolate/Mutable/src/Types.Mutable/Serialization/SchemaParser.cs b/src/HotChocolate/Mutable/src/Types.Mutable/Serialization/SchemaParser.cs

Original file line number	Diff line number	Diff line change
`@@ -5,11 +5,6 @@ namespace HotChocolate.Fusion.Extensions;`
`5`	`5`
`6`	`6`	`internal static class MutableObjectTypeDefinitionExtensions`
`7`	`7`	`{`
`8`		`- public static void ApplyShareableDirective(this MutableObjectTypeDefinition type)`
`9`		`- {`
`10`		`- type.Directives.Add(new Directive(new MutableDirectiveDefinition(DirectiveNames.Shareable)));`
`11`		`- }`
`12`		`-`
`13`	`8`	`public static bool HasInternalDirective(this MutableObjectTypeDefinition type)`
`14`	`9`	`{`
`15`	`10`	`return type.Directives.ContainsName(DirectiveNames.Internal);`
Original file line number	Diff line number	Diff line change
`@@ -9,15 +9,15 @@ internal static class MutableSchemaDefinitionExtensions`
`9`	`9`	`{`
`10`	`10`	`public static void AddBuiltInFusionTypes(this MutableSchemaDefinition schema)`
`11`	`11`	`{`
`12`		`- foreach (var builtInScalar in FusionBuiltIns.SourceSchemaScalars)`
	`12`	`+ foreach (var builtInScalar in FusionBuiltIns.SourceSchemaScalars.Values)`
`13`	`13`	`{`
`14`	`14`	`schema.Types.Add(builtInScalar);`
`15`	`15`	`}`
`16`	`16`	`}`
`17`	`17`
`18`	`18`	`public static void AddBuiltInFusionDirectives(this MutableSchemaDefinition schema)`
`19`	`19`	`{`
`20`		`- foreach (var builtInDirective in FusionBuiltIns.SourceSchemaDirectives)`
	`20`	`+ foreach (var builtInDirective in FusionBuiltIns.SourceSchemaDirectives.Values)`
`21`	`21`	`{`
`22`	`22`	`schema.DirectiveDefinitions.Add(builtInDirective);`
`23`	`23`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3,4 +3,6 @@ namespace HotChocolate.Fusion.Options;`
`3`	`3`	`public sealed class SchemaComposerOptions`
`4`	`4`	`{`
`5`	`5`	`public required bool EnableGlobalObjectIdentification { get; init; }`
	`6`	`+`
	`7`	`+ public SourceSchemaPreprocessorOptions Preprocessor { get; init; } = new();`
`6`	`8`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,9 @@`
`1`	`1`	`namespace HotChocolate.Fusion.Options;`
`2`	`2`
`3`		`-internal sealed class SourceSchemaPreprocessorOptions`
	`3`	`+public sealed class SourceSchemaPreprocessorOptions`
`4`	`4`	`{`
`5`		`- public bool ApplyShareableToAllTypes { get; set; } = true;`
	`5`	`+ /// <summary>`
	`6`	`+ /// Applies inferred key directives to types that are returned by lookup fields.`
	`7`	`+ /// </summary>`
	`8`	`+ public bool ApplyInferredKeyDirectives { get; set; } = true;`
`6`	`9`	`}`