Added support for reference resolvers on an entity interface.

Author: danielreynolds1

src/HotChocolate/ApolloFederation/src/ApolloFederation/FederationTypeInterceptor.cs

@@ -36,6 +36,7 @@ internal sealed class FederationTypeInterceptor : TypeInterceptor
                 BindingFlags.Static | BindingFlags.Public)!;
 
     private readonly List<ObjectType> _entityTypes = [];
+    private readonly List<InterfaceType> _entityInterfaceTypes = [];
     private readonly Dictionary<Uri, HashSet<string>> _imports = [];
     private IDescriptorContext _context = null!;
     private ITypeInspector _typeInspector = null!;
@@ -79,6 +80,25 @@ public override void OnAfterInitialize(
                 objectType,
                 objectTypeCfg,
                 discoveryContext);
+            return;
+        }
+
+        if (discoveryContext.Type is InterfaceType interfaceType &&
+            configuration is InterfaceTypeConfiguration interfaceTypeCfg)
+        {
+            ApplyMethodLevelReferenceResolvers(
+                interfaceType,
+                interfaceTypeCfg);
+
+            AddToUnionIfHasTypeLevelKeyDirective(
+                interfaceType,
+                interfaceTypeCfg);
+
+            AggregatePropertyLevelKeyDirectives(
+                interfaceType,
+                interfaceTypeCfg,
+                discoveryContext);
+            return;
         }
     }
 
@@ -345,12 +365,24 @@ public override void OnAfterMakeExecutable(
         {
             CompleteExternalFieldSetters(type, typeCfg);
             CompleteReferenceResolver(typeCfg);
+            return;
+        }
+
+        if (completionContext.Type is InterfaceType interfaceType &&
+            configuration is InterfaceTypeConfiguration interfaceTypeCfg)
+        {
+            CompleteExternalFieldSetters(interfaceType, interfaceTypeCfg);
+            CompleteReferenceResolver(interfaceTypeCfg);
+            return;
         }
     }
 
     private void CompleteExternalFieldSetters(ObjectType type, ObjectTypeConfiguration typeCfg)
         => ExternalSetterExpressionHelper.TryAddExternalSetter(type, typeCfg);
 
+    private void CompleteExternalFieldSetters(InterfaceType type, InterfaceTypeConfiguration typeCfg)
+        => ExternalSetterExpressionHelper.TryAddExternalSetter(type, typeCfg);
+
     private void CompleteReferenceResolver(ObjectTypeConfiguration typeCfg)
     {
         var resolvers = typeCfg.Features.Get<List<ReferenceResolverConfiguration>>();
@@ -398,6 +430,53 @@ private void CompleteReferenceResolver(ObjectTypeConfiguration typeCfg)
         }
     }
 
+    private void CompleteReferenceResolver(InterfaceTypeConfiguration typeCfg)
+    {
+        var resolvers = typeCfg.Features.Get<List<ReferenceResolverConfiguration>>();
+
+        if (resolvers is null)
+        {
+            return;
+        }
+
+        if (resolvers.Count == 1)
+        {
+            typeCfg.Features.Set(new ReferenceResolver(resolvers[0].Resolver));
+        }
+        else
+        {
+            var expressions = new Stack<(Expression Condition, Expression Execute)>();
+            var context = Expression.Parameter(typeof(IResolverContext));
+
+            foreach (var resolverDef in resolvers)
+            {
+                Expression required = Expression.Constant(resolverDef.Required);
+                Expression resolver = Expression.Constant(resolverDef.Resolver);
+                Expression condition = Expression.Call(s_matches, context, required);
+                Expression execute = Expression.Call(s_execute, context, resolver);
+                expressions.Push((condition, execute));
+            }
+
+            Expression current = Expression.Call(s_invalid, context);
+            var variable = Expression.Variable(typeof(ValueTask<object?>));
+
+            while (expressions.Count > 0)
+            {
+                var expression = expressions.Pop();
+                current = Expression.IfThenElse(
+                    expression.Condition,
+                    Expression.Assign(variable, expression.Execute),
+                    current);
+            }
+
+            current = Expression.Block([variable], current, variable);
+
+            typeCfg.Features.Set(
+                new ReferenceResolver(
+                    Expression.Lambda<FieldResolverDelegate>(current, context).Compile()));
+        }
+    }
+
     private void AddServiceTypeToQueryType(
         ITypeCompletionContext completionContext,
         TypeSystemConfiguration? definition)
@@ -409,7 +488,7 @@ private void AddServiceTypeToQueryType(
 
         var objectTypeCfg = (ObjectTypeConfiguration)definition!;
         objectTypeCfg.Fields.Add(ServerFields.CreateServiceField(_context));
-        if (_entityTypes.Count > 0)
+        if (_entityTypes.Count > 0 || _entityInterfaceTypes.Count > 0)
         {
             objectTypeCfg.Fields.Add(ServerFields.CreateEntitiesField(_context));
         }
@@ -452,6 +531,43 @@ private void ApplyMethodLevelReferenceResolvers(
         descriptor.CreateConfiguration();
     }
 
+    private void ApplyMethodLevelReferenceResolvers(
+        InterfaceType interfaceType,
+        InterfaceTypeConfiguration interfaceTypeCfg)
+    {
+        if (interfaceType.RuntimeType == typeof(object))
+        {
+            return;
+        }
+
+        var descriptor = InterfaceTypeDescriptor.From(_context, interfaceTypeCfg);
+
+        // Static methods won't end up in the schema as fields.
+        // The default initialization system only considers instance methods,
+        // so we have to handle the attributes for those manually.
+        var potentiallyUnregisteredReferenceResolvers = interfaceType.RuntimeType
+            .GetMethods(BindingFlags.Static | BindingFlags.Public);
+
+        foreach (var possibleReferenceResolver in potentiallyUnregisteredReferenceResolvers)
+        {
+            if (!possibleReferenceResolver.IsDefined(typeof(ReferenceResolverAttribute)))
+            {
+                continue;
+            }
+
+            foreach (var attribute in possibleReferenceResolver.GetCustomAttributes(true))
+            {
+                if (attribute is ReferenceResolverAttribute casted)
+                {
+                    casted.TryConfigure(_context, descriptor, possibleReferenceResolver);
+                }
+            }
+        }
+
+        // This seems to re-detect the entity resolver and save it into the context data.
+        descriptor.CreateConfiguration();
+    }
+
     private void AddToUnionIfHasTypeLevelKeyDirective(
         ObjectType objectType,
         ObjectTypeConfiguration objectTypeCfg)
@@ -469,6 +585,23 @@ private void AddToUnionIfHasTypeLevelKeyDirective(
         }
     }
 
+    private void AddToUnionIfHasTypeLevelKeyDirective(
+        InterfaceType interfaceType,
+        InterfaceTypeConfiguration interfaceTypeCfg)
+    {
+        if (interfaceTypeCfg.Directives.FirstOrDefault(d => d.Value is KeyDirective) is { } keyDirective &&
+            ((KeyDirective)keyDirective.Value).Resolvable)
+        {
+            _entityInterfaceTypes.Add(interfaceType);
+            return;
+        }
+
+        if (interfaceTypeCfg.Fields.Any(f => f.Features.TryGet(out KeyMarker? key) && key.Resolvable))
+        {
+            _entityInterfaceTypes.Add(interfaceType);
+        }
+    }
+
     private void AggregatePropertyLevelKeyDirectives(
         ObjectType objectType,
         ObjectTypeConfiguration objectTypeCfg,
@@ -532,6 +665,69 @@ private void AggregatePropertyLevelKeyDirectives(
         }
     }
 
+    private void AggregatePropertyLevelKeyDirectives(
+        InterfaceType interfaceType,
+        InterfaceTypeConfiguration interfaceTypeCfg,
+        ITypeDiscoveryContext discoveryContext)
+    {
+        // if we find key markers on our fields, we need to construct the key directive
+        // from the annotated fields.
+        var foundMarkers = interfaceTypeCfg.Fields.Any(f => f.Features.TryGet(out KeyMarker? _));
+
+        if (!foundMarkers)
+        {
+            return;
+        }
+
+        IReadOnlyList<InterfaceFieldConfiguration> fields = interfaceTypeCfg.Fields;
+        var fieldSet = new StringBuilder();
+        bool? resolvable = null;
+
+        foreach (var fieldDefinition in fields)
+        {
+            if (fieldDefinition.Features.TryGet(out KeyMarker? key))
+            {
+                if (resolvable is null)
+                {
+                    resolvable = key.Resolvable;
+                }
+                else if (resolvable != key.Resolvable)
+                {
+                    throw Key_FieldSet_ResolvableMustBeConsistent(fieldDefinition.Member!);
+                }
+
+                if (fieldSet.Length > 0)
+                {
+                    fieldSet.Append(' ');
+                }
+
+                fieldSet.Append(fieldDefinition.Name);
+            }
+        }
+
+        // add the key directive with the dynamically generated field set.
+        AddKeyDirective(interfaceTypeCfg, fieldSet.ToString(), resolvable ?? true);
+
+        // register dependency to the key directive so that it is completed before
+        // we complete this type.
+        foreach (var directiveDefinition in interfaceTypeCfg.Directives)
+        {
+            discoveryContext.Dependencies.Add(
+                new TypeDependency(
+                    directiveDefinition.Type,
+                    TypeDependencyFulfilled.Completed));
+
+            discoveryContext.Dependencies.Add(new(directiveDefinition.Type));
+        }
+
+        if (resolvable ?? true)
+        {
+            // since this type has now a key directive we also need to add this type to
+            // the _Entity union type provided that the key is resolvable.
+            _entityInterfaceTypes.Add(interfaceType);
+        }
+    }
+
     private void AddMemberTypesToTheEntityUnionType(
         ITypeCompletionContext completionContext,
         TypeSystemConfiguration? definition)
@@ -556,4 +752,15 @@ private void AddKeyDirective(
                 new KeyDirective(fieldSet, resolvable),
                 _keyDirectiveReference));
     }
+
+    private void AddKeyDirective(
+        InterfaceTypeConfiguration interfaceTypeCfg,
+        string fieldSet,
+        bool resolvable)
+    {
+        interfaceTypeCfg.Directives.Add(
+            new DirectiveConfiguration(
+                new KeyDirective(fieldSet, resolvable),
+                _keyDirectiveReference));
+    }
 }

src/HotChocolate/ApolloFederation/src/ApolloFederation/Resolvers/EntitiesResolver.cs

@@ -34,11 +34,24 @@ internal static class EntitiesResolver
                 entityContext.SetLocalState(DataField, current.Data);
 
                 tasks[i] = entity.Resolver.Invoke(entityContext).AsTask();
+                continue;
             }
-            else
+
+            if (schema.Types.TryGetType<InterfaceType>(current.TypeName, out var interfaceType) &&
+                interfaceType.Features.TryGet(out ReferenceResolver? entityInterface))
             {
-                throw ThrowHelper.EntityResolver_NoResolverFound();
+                // We clone the resolver context here so that we can split the work
+                // into subtasks that can be awaited in parallel and produce separate results.
+                var entityContext = context.Clone();
+
+                entityContext.SetLocalState(TypeField, interfaceType);
+                entityContext.SetLocalState(DataField, current.Data);
+
+                tasks[i] = entityInterface.Resolver.Invoke(entityContext).AsTask();
+                continue;
             }
+
+            throw ThrowHelper.EntityResolver_NoResolverFound();
         }
 
         for (var i = 0; i < representations.Count; i++)

src/HotChocolate/ApolloFederation/src/ApolloFederation/Resolvers/ExternalSetterExpressionHelper.cs

@@ -47,6 +47,30 @@ public static void TryAddExternalSetter(ObjectType type, ObjectTypeConfiguration
         }
     }
 
+    public static void TryAddExternalSetter(InterfaceType type, InterfaceTypeConfiguration typeDef)
+    {
+        List<Expression>? block = null;
+
+        foreach (var field in type.Fields)
+        {
+            if (field.Directives.ContainsDirective<ExternalDirective>() &&
+                typeDef.Fields.FirstOrDefault(f => f.Name == field.Name) is
+                    { Member: PropertyInfo { SetMethod: not null } property })
+            {
+                var expression = CreateTrySetValue(type.RuntimeType, property, field.Name);
+                (block ??= []).Add(expression);
+            }
+        }
+
+        if (block is not null)
+        {
+            typeDef.Features.Set(new ExternalSetter(
+                Lambda<Action<ObjectType, IValueNode, object>>(
+                        Block(block), s_type, s_data, s_entity)
+                    .Compile()));
+        }
+    }
+
     private static Expression CreateTrySetValue(
         Type runtimeType,
         PropertyInfo property,

src/HotChocolate/ApolloFederation/src/ApolloFederation/Resolvers/ReferenceResolverArgumentExpressionBuilder.cs

@@ -14,6 +14,14 @@ internal sealed class ReferenceResolverArgumentExpressionBuilder :
             nameof(ArgumentParser.GetValue),
             BindingFlags.Static | BindingFlags.Public)!;
 
+    private readonly Type _targetType;
+
+    public ReferenceResolverArgumentExpressionBuilder(Type targetType)
+    {
+        ArgumentNullException.ThrowIfNull(targetType);
+        _targetType = targetType;
+    }
+
     public override Expression Build(ParameterExpressionBuilderContext context)
     {
         var param = context.Parameter;
@@ -35,7 +43,7 @@ public override Expression Build(ParameterExpressionBuilderContext context)
             param,
             typeKey,
             context.ResolverContext,
-            typeof(ObjectType));
+            _targetType);
         var getValueMethod = _getValue.MakeGenericMethod(param.ParameterType);
         var getValue = Expression.Call(
             getValueMethod,