From bdf43313871d8132e3660432891e9ddcc6ce42e9 Mon Sep 17 00:00:00 2001
From: Eddy Nguyen <github@eddeee888.me>
Date: Tue, 10 Dec 2024 23:40:09 +1100
Subject: [PATCH] Fix reference being assigned mappers incorrectly

---
 .../src/base-resolvers-visitor.ts             | 34 ++++++++++++++++++-
 .../plugins/typescript/resolvers/src/index.ts |  2 ++
 .../utils/plugins-helpers/src/federation.ts   |  6 ++--
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/packages/plugins/other/visitor-plugin-common/src/base-resolvers-visitor.ts b/packages/plugins/other/visitor-plugin-common/src/base-resolvers-visitor.ts
index 19351915e15..b82049905ef 100644
--- a/packages/plugins/other/visitor-plugin-common/src/base-resolvers-visitor.ts
+++ b/packages/plugins/other/visitor-plugin-common/src/base-resolvers-visitor.ts
@@ -1223,6 +1223,28 @@ export class BaseResolversVisitor<
       ).string;
   }
 
+  public buildFederationTypes(): string {
+    const federationMeta = this._federation.getMeta();
+
+    if (Object.keys(federationMeta).length === 0) {
+      return '';
+    }
+
+    const declarationKind = 'type';
+    return new DeclarationBlock(this._declarationBlockConfig)
+      .export()
+      .asKind(declarationKind)
+      .withName(this.convertName('FederationTypes'))
+      .withComment('Mapping of federation types')
+      .withBlock(
+        Object.keys(federationMeta)
+          .map(typeName => {
+            return indent(`${typeName}: ${this.convertName(typeName)}${this.getPunctuation(declarationKind)}`);
+          })
+          .join('\n')
+      ).string;
+  }
+
   public get schema(): GraphQLSchema {
     return this._schema;
   }
@@ -1498,6 +1520,7 @@ export class BaseResolversVisitor<
         fieldNode: original,
         parentType,
         parentTypeSignature: this.getParentTypeForSignature(node),
+        federationTypeSignature: 'FederationType',
       });
       const mappedTypeKey = isSubscriptionType ? `${mappedType}, "${node.name}"` : mappedType;
 
@@ -1619,10 +1642,19 @@ export class BaseResolversVisitor<
       );
     }
 
+    const genericTypes: string[] = [
+      `ContextType = ${this.config.contextType.type}`,
+      this.transformParentGenericType(parentType),
+    ];
+    if (this._federation.getMeta()[typeName]) {
+      const typeRef = `${this.convertName('FederationTypes')}['${typeName}']`;
+      genericTypes.push(`FederationType extends ${typeRef} = ${typeRef}`);
+    }
+
     const block = new DeclarationBlock(this._declarationBlockConfig)
       .export()
       .asKind(declarationKind)
-      .withName(name, `<ContextType = ${this.config.contextType.type}, ${this.transformParentGenericType(parentType)}>`)
+      .withName(name, `<${genericTypes.join(', ')}>`)
       .withBlock(fieldsContent.join('\n'));
 
     this._collectedResolvers[node.name as any] = {
diff --git a/packages/plugins/typescript/resolvers/src/index.ts b/packages/plugins/typescript/resolvers/src/index.ts
index e0d3ff293a3..e949ae972a2 100644
--- a/packages/plugins/typescript/resolvers/src/index.ts
+++ b/packages/plugins/typescript/resolvers/src/index.ts
@@ -244,6 +244,7 @@ export type DirectiveResolverFn<TResult = {}, TParent = {}, TContext = {}, TArgs
 ) => TResult | Promise<TResult>;
 `;
 
+  const federationTypes = visitor.buildFederationTypes();
   const resolversTypeMapping = visitor.buildResolversTypes();
   const resolversParentTypeMapping = visitor.buildResolversParentTypes();
   const resolversUnionTypesMapping = visitor.buildResolversUnionTypes();
@@ -287,6 +288,7 @@ export type DirectiveResolverFn<TResult = {}, TParent = {}, TContext = {}, TArgs
     prepend,
     content: [
       header,
+      federationTypes,
       resolversUnionTypesMapping,
       resolversInterfaceTypesMapping,
       resolversTypeMapping,
diff --git a/packages/utils/plugins-helpers/src/federation.ts b/packages/utils/plugins-helpers/src/federation.ts
index d77277629e7..685150e6c10 100644
--- a/packages/utils/plugins-helpers/src/federation.ts
+++ b/packages/utils/plugins-helpers/src/federation.ts
@@ -155,10 +155,12 @@ export class ApolloFederation {
     fieldNode,
     parentType,
     parentTypeSignature,
+    federationTypeSignature,
   }: {
     fieldNode: FieldDefinitionNode;
     parentType: GraphQLNamedType;
     parentTypeSignature: string;
+    federationTypeSignature: string;
   }) {
     if (
       this.enabled &&
@@ -177,12 +179,12 @@ export class ApolloFederation {
 
         // Look for @requires and see what the service needs and gets
         const requires = getDirectivesByName('requires', fieldNode).map(this.extractFieldSet);
-        const requiredFields = this.translateFieldSet(merge({}, ...requires), parentTypeSignature);
+        const requiredFields = this.translateFieldSet(merge({}, ...requires), federationTypeSignature);
 
         // @key() @key() - "primary keys" in Federation
         const primaryKeys = resolvableKeyDirectives.map(def => {
           const fields = this.extractFieldSet(def);
-          return this.translateFieldSet(fields, parentTypeSignature);
+          return this.translateFieldSet(fields, federationTypeSignature);
         });
 
         const [open, close] = primaryKeys.length > 1 ? ['(', ')'] : ['', ''];