[typescript-resolvers] Extract union types to ResolversUnionTypes (#9069)

Author: eddeee888

.changeset/brave-wasps-tap.md

@@ -0,0 +1,6 @@
+---
+'@graphql-codegen/visitor-plugin-common': patch
+'@graphql-codegen/typescript-resolvers': patch
+---
+
+Extract union types to ResolversUnionTypes

dev-test/modules/types.ts

@@ -162,6 +162,11 @@ export type DirectiveResolverFn<TResult = {}, TParent = {}, TContext = {}, TArgs
   info: GraphQLResolveInfo
 ) => TResult | Promise<TResult>;
 
+/** Mapping of union types */
+export type ResolversUnionTypes = {
+  PaymentOption: CreditCard | Paypal;
+};
+
 /** Mapping between all available schema types and the resolvers types */
 export type ResolversTypes = {
   Article: ResolverTypeWrapper<Article>;
@@ -173,7 +178,7 @@ export type ResolversTypes = {
   ID: ResolverTypeWrapper<Scalars['ID']>;
   Int: ResolverTypeWrapper<Scalars['Int']>;
   Mutation: ResolverTypeWrapper<{}>;
-  PaymentOption: ResolversTypes['CreditCard'] | ResolversTypes['Paypal'];
+  PaymentOption: ResolverTypeWrapper<ResolversUnionTypes['PaymentOption']>;
   Paypal: ResolverTypeWrapper<Paypal>;
   Query: ResolverTypeWrapper<{}>;
   String: ResolverTypeWrapper<Scalars['String']>;
@@ -193,7 +198,7 @@ export type ResolversParentTypes = {
   ID: Scalars['ID'];
   Int: Scalars['Int'];
   Mutation: {};
-  PaymentOption: ResolversParentTypes['CreditCard'] | ResolversParentTypes['Paypal'];
+  PaymentOption: ResolversUnionTypes['PaymentOption'];
   Paypal: Paypal;
   Query: {};
   String: Scalars['String'];

packages/plugins/other/visitor-plugin-common/src/base-resolvers-visitor.ts

@@ -560,6 +560,8 @@ export class BaseResolversVisitor<
   protected _usedMappers: { [key: string]: boolean } = {};
   protected _resolversTypes: ResolverTypes = {};
   protected _resolversParentTypes: ResolverParentTypes = {};
+  protected _hasReferencedResolversUnionTypes = false;
+  protected _resolversUnionTypes: Record<string, string> = {};
   protected _rootTypeNames = new Set<string>();
   protected _globalDeclarations = new Set<string>();
   protected _federation: ApolloFederation;
@@ -624,6 +626,7 @@ export class BaseResolversVisitor<
       name => this.getParentTypeToUse(name),
       namedType => !isEnumType(namedType)
     );
+    this._resolversUnionTypes = this.createResolversUnionTypes();
     this._fieldContextTypeMap = this.createFieldContextTypeMap();
     this._directiveContextTypesMap = this.createDirectivedContextType();
     this._directiveResolverMappings = rawConfig.directiveResolverMappings ?? {};
@@ -754,10 +757,9 @@ export class BaseResolversVisitor<
       } else if (isScalar) {
         prev[typeName] = applyWrapper(this._getScalar(typeName));
       } else if (isUnionType(schemaType)) {
-        prev[typeName] = schemaType
-          .getTypes()
-          .map(type => getTypeToUse(type.name))
-          .join(' | ');
+        this._hasReferencedResolversUnionTypes = true;
+        const resolversType = this.convertName('ResolversUnionTypes');
+        prev[typeName] = applyWrapper(`${resolversType}['${typeName}']`);
       } else if (isEnumType(schemaType)) {
         prev[typeName] = this.convertName(typeName, { useTypesPrefix: this.config.enumPrefix }, true);
       } else {
@@ -766,44 +768,11 @@ export class BaseResolversVisitor<
       }
 
       if (shouldApplyOmit && prev[typeName] !== 'any' && isObjectType(schemaType)) {
-        const fields = schemaType.getFields();
-        const relevantFields: {
-          addOptionalSign: boolean;
-          fieldName: string;
-          replaceWithType: string;
-        }[] = this._federation
-          .filterFieldNames(Object.keys(fields))
-          .filter(fieldName => {
-            const field = fields[fieldName];
-            const baseType = getBaseType(field.type);
-
-            // Filter out fields of types that are not included
-            if (shouldInclude && !shouldInclude(baseType)) {
-              return false;
-            }
-            return true;
-          })
-          .map(fieldName => {
-            const field = fields[fieldName];
-            const baseType = getBaseType(field.type);
-            const isUnion = isUnionType(baseType);
-
-            if (!this.config.mappers[baseType.name] && !isUnion && !this._shouldMapType[baseType.name]) {
-              return null;
-            }
-
-            const addOptionalSign = !this.config.avoidOptionals && !isNonNullType(field.type);
-
-            return {
-              addOptionalSign,
-              fieldName,
-              replaceWithType: wrapTypeWithModifiers(getTypeToUse(baseType.name), field.type, {
-                wrapOptional: this.applyMaybe,
-                wrapArray: this.wrapWithArray,
-              }),
-            };
-          })
-          .filter(a => a);
+        const relevantFields = this.getRelevantFieldsToOmit({
+          schemaType,
+          getTypeToUse,
+          shouldInclude,
+        });
 
         if (relevantFields.length > 0) {
           // Puts ResolverTypeWrapper on top of an entire type
@@ -819,15 +788,15 @@ export class BaseResolversVisitor<
       }
 
       if (!isMapped && hasDefaultMapper && hasPlaceholder(this.config.defaultMapper.type)) {
-        // Make sure the inner type has no ResolverTypeWrapper
-        const name = clearWrapper(isScalar ? this._getScalar(typeName) : prev[typeName]);
-        const replaced = replacePlaceholder(this.config.defaultMapper.type, name);
+        const originalTypeName = isScalar ? this._getScalar(typeName) : prev[typeName];
 
-        // Don't wrap Union with ResolverTypeWrapper, each inner type already has it
         if (isUnionType(schemaType)) {
-          prev[typeName] = replaced;
+          // Don't clear ResolverTypeWrapper from Unions
+          prev[typeName] = replacePlaceholder(this.config.defaultMapper.type, originalTypeName);
         } else {
-          prev[typeName] = applyWrapper(replacePlaceholder(this.config.defaultMapper.type, name));
+          const name = clearWrapper(originalTypeName);
+          const replaced = replacePlaceholder(this.config.defaultMapper.type, name);
+          prev[typeName] = applyWrapper(replaced);
         }
       }
 
@@ -837,7 +806,7 @@ export class BaseResolversVisitor<
 
   protected replaceFieldsInType(
     typeName: string,
-    relevantFields: { addOptionalSign: boolean; fieldName: string; replaceWithType: string }[]
+    relevantFields: ReturnType<typeof this.getRelevantFieldsToOmit>
   ): string {
     this._globalDeclarations.add(OMIT_TYPE);
     return `Omit<${typeName}, ${relevantFields.map(f => `'${f.fieldName}'`).join(' | ')}> & { ${relevantFields
@@ -880,6 +849,64 @@ export class BaseResolversVisitor<
     return `Array<${t}>`;
   }
 
+  protected createResolversUnionTypes(): Record<string, string> {
+    if (!this._hasReferencedResolversUnionTypes) {
+      return {};
+    }
+
+    const allSchemaTypes = this._schema.getTypeMap();
+    const typeNames = this._federation.filterTypeNames(Object.keys(allSchemaTypes));
+
+    const unionTypes = typeNames.reduce((res, typeName) => {
+      const schemaType = allSchemaTypes[typeName];
+
+      if (isUnionType(schemaType)) {
+        const referencedTypes = schemaType.getTypes().map(unionMemberType => {
+          const isUnionMemberMapped = this.config.mappers[unionMemberType.name];
+
+          // 1. If mapped without placehoder, just use it without doing extra checks
+          if (isUnionMemberMapped && !hasPlaceholder(isUnionMemberMapped.type)) {
+            return isUnionMemberMapped.type;
+          }
+
+          // 2. Work out value for union member type
+          // 2a. By default, use the typescript type
+          let unionMemberValue = this.convertName(unionMemberType.name, {}, true);
+
+          // 2b. Find fields to Omit if needed.
+          //  - If no field to Omit, "type with maybe Omit" is typescript type i.e. no Omit
+          //  - If there are fields to Omit, "type with maybe Omit"
+          const fieldsToOmit = this.getRelevantFieldsToOmit({
+            schemaType: unionMemberType,
+            getTypeToUse: this.getTypeToUse,
+          });
+          if (fieldsToOmit.length > 0) {
+            unionMemberValue = this.replaceFieldsInType(unionMemberValue, fieldsToOmit);
+          }
+
+          // 2c. If union member is mapped with placeholder, use the "type with maybe Omit" as {T}
+          if (isUnionMemberMapped && hasPlaceholder(isUnionMemberMapped.type)) {
+            return replacePlaceholder(isUnionMemberMapped.type, unionMemberValue);
+          }
+
+          // 2d. If has default mapper with placeholder, use the "type with maybe Omit" as {T}
+          const hasDefaultMapper = !!this.config.defaultMapper?.type;
+          const isScalar = this.config.scalars[typeName];
+          if (hasDefaultMapper && hasPlaceholder(this.config.defaultMapper.type)) {
+            const finalTypename = isScalar ? this._getScalar(typeName) : unionMemberValue;
+            return replacePlaceholder(this.config.defaultMapper.type, finalTypename);
+          }
+
+          return unionMemberValue;
+        });
+        res[typeName] = referencedTypes.map(type => `( ${type} )`).join(' | '); // Must wrap every union member in explicit "( )" to separate the members
+      }
+      return res;
+    }, {});
+
+    return unionTypes;
+  }
+
   protected createFieldContextTypeMap(): FieldContextTypeMap {
     return this.config.fieldContextTypes.reduce<FieldContextTypeMap>((prev, fieldContextType) => {
       const items = fieldContextType.split('#');
@@ -935,6 +962,24 @@ export class BaseResolversVisitor<
       ).string;
   }
 
+  public buildResolversUnionTypes(): string {
+    if (Object.keys(this._resolversUnionTypes).length === 0) {
+      return '';
+    }
+
+    const declarationKind = 'type';
+    return new DeclarationBlock(this._declarationBlockConfig)
+      .export()
+      .asKind(declarationKind)
+      .withName(this.convertName('ResolversUnionTypes'))
+      .withComment('Mapping of union types')
+      .withBlock(
+        Object.entries(this._resolversUnionTypes)
+          .map(([typeName, value]) => indent(`${typeName}: ${value}${this.getPunctuation(declarationKind)}`))
+          .join('\n')
+      ).string;
+  }
+
   public get schema(): GraphQLSchema {
     return this._schema;
   }
@@ -1478,6 +1523,55 @@ export class BaseResolversVisitor<
   SchemaDefinition() {
     return null;
   }
+
+  private getRelevantFieldsToOmit({
+    schemaType,
+    shouldInclude,
+    getTypeToUse,
+  }: {
+    schemaType: GraphQLObjectType;
+    getTypeToUse: (name: string) => string;
+    shouldInclude?: (type: GraphQLNamedType) => boolean;
+  }): {
+    addOptionalSign: boolean;
+    fieldName: string;
+    replaceWithType: string;
+  }[] {
+    const fields = schemaType.getFields();
+    return this._federation
+      .filterFieldNames(Object.keys(fields))
+      .filter(fieldName => {
+        const field = fields[fieldName];
+        const baseType = getBaseType(field.type);
+
+        // Filter out fields of types that are not included
+        if (shouldInclude && !shouldInclude(baseType)) {
+          return false;
+        }
+        return true;
+      })
+      .map(fieldName => {
+        const field = fields[fieldName];
+        const baseType = getBaseType(field.type);
+        const isUnion = isUnionType(baseType);
+
+        if (!this.config.mappers[baseType.name] && !isUnion && !this._shouldMapType[baseType.name]) {
+          return null;
+        }
+
+        const addOptionalSign = !this.config.avoidOptionals && !isNonNullType(field.type);
+
+        return {
+          addOptionalSign,
+          fieldName,
+          replaceWithType: wrapTypeWithModifiers(getTypeToUse(baseType.name), field.type, {
+            wrapOptional: this.applyMaybe,
+            wrapArray: this.wrapWithArray,
+          }),
+        };
+      })
+      .filter(a => a);
+  }
 }
 
 function replacePlaceholder(pattern: string, typename: string): string {

packages/plugins/typescript/resolvers/src/index.ts

@@ -245,6 +245,7 @@ export type DirectiveResolverFn<TResult = {}, TParent = {}, TContext = {}, TArgs
 
   const resolversTypeMapping = visitor.buildResolversTypes();
   const resolversParentTypeMapping = visitor.buildResolversParentTypes();
+  const resolversUnionTypesMapping = visitor.buildResolversUnionTypes();
   const { getRootResolver, getAllDirectiveResolvers, mappersImports, unusedMappers, hasScalars } = visitor;
 
   if (hasScalars()) {
@@ -282,6 +283,7 @@ export type DirectiveResolverFn<TResult = {}, TParent = {}, TContext = {}, TArgs
     prepend,
     content: [
       header,
+      resolversUnionTypesMapping,
       resolversTypeMapping,
       resolversParentTypeMapping,
       ...visitorResult.definitions.filter(d => typeof d === 'string'),