enhance(stitch/federation): improvements on field merging and extraction of unavailable fields

ardatan · ardatan · commit 9238e140862d · 2024-05-01T17:32:13.000+03:00
diff --git a/.changeset/soft-otters-mix.md b/.changeset/soft-otters-mix.md
@@ -0,0 +1,6 @@
+---
+"@graphql-tools/federation": patch
+"@graphql-tools/stitch": patch
+---
+
+Improvements on field merging and extraction of unavailable fields
diff --git a/packages/federation/src/supergraph.ts b/packages/federation/src/supergraph.ts
@@ -4,6 +4,7 @@ import {
   EnumTypeDefinitionNode,
   EnumValueDefinitionNode,
   FieldDefinitionNode,
+  GraphQLOutputType,
   GraphQLSchema,
   InputValueDefinitionNode,
   InterfaceTypeDefinitionNode,
@@ -22,8 +23,14 @@ import {
 } from 'graphql';
 import { MergedTypeConfig, SubschemaConfig } from '@graphql-tools/delegate';
 import { buildHTTPExecutor } from '@graphql-tools/executor-http';
-import { stitchSchemas } from '@graphql-tools/stitch';
-import { type Executor } from '@graphql-tools/utils';
+import {
+  getDefaultFieldConfigMerger,
+  MergeFieldConfigCandidate,
+  stitchSchemas,
+  TypeMergingOptions,
+  ValidationLevel,
+} from '@graphql-tools/stitch';
+import { memoize1, type Executor } from '@graphql-tools/utils';
 import {
   filterInternalFieldsAndTypes,
   getArgsFromKeysForFederation,
@@ -42,13 +49,60 @@ export interface GetSubschemasFromSupergraphSdlOpts {
   batch?: boolean;
 }
 
+export function ensureSupergraphSDLAst(supergraphSdl: string | DocumentNode): DocumentNode {
+  return typeof supergraphSdl === 'string'
+    ? parse(supergraphSdl, { noLocation: true })
+    : supergraphSdl;
+}
+
+function getTypeFieldMapFromSupergraphAST(supergraphAST: DocumentNode) {
+  const typeFieldASTMap = new Map<
+    string,
+    Map<string, FieldDefinitionNode | InputValueDefinitionNode>
+  >();
+  for (const definition of supergraphAST.definitions) {
+    if ('fields' in definition) {
+      const fieldMap = new Map<string, FieldDefinitionNode | InputValueDefinitionNode>();
+      typeFieldASTMap.set(definition.name.value, fieldMap);
+      for (const field of definition.fields || []) {
+        fieldMap.set(field.name.value, field);
+      }
+    }
+  }
+  return typeFieldASTMap;
+}
+
+export function getFieldMergerFromSupergraphSdl(
+  supergraphSdl: DocumentNode | string,
+): TypeMergingOptions['fieldConfigMerger'] {
+  const supergraphAST = ensureSupergraphSDLAst(supergraphSdl);
+  const typeFieldASTMap = getTypeFieldMapFromSupergraphAST(supergraphAST);
+  const defaultMerger = getDefaultFieldConfigMerger(true);
+  const memoizedASTPrint = memoize1(print);
+  const memoizedTypePrint = memoize1((type: GraphQLOutputType) => type.toString());
+  return function (candidates: MergeFieldConfigCandidate[]) {
+    const filteredCandidates = candidates.filter(candidate => {
+      const fieldASTMap = typeFieldASTMap.get(candidate.type.name);
+      if (fieldASTMap) {
+        const fieldAST = fieldASTMap.get(candidate.fieldName);
+        if (fieldAST) {
+          const typeNodeInAST = memoizedASTPrint(fieldAST.type);
+          const typeNodeInCandidate = memoizedTypePrint(candidate.fieldConfig.type);
+          return typeNodeInAST === typeNodeInCandidate;
+        }
+      }
+      return false;
+    });
+    return defaultMerger(filteredCandidates.length ? filteredCandidates : candidates);
+  };
+}
+
 export function getSubschemasFromSupergraphSdl({
   supergraphSdl,
   onExecutor = ({ endpoint }) => buildHTTPExecutor({ endpoint }),
   batch = false,
 }: GetSubschemasFromSupergraphSdlOpts) {
-  const ast =
-    typeof supergraphSdl === 'string' ? parse(supergraphSdl, { noLocation: true }) : supergraphSdl;
+  const supergraphAst = ensureSupergraphSDLAst(supergraphSdl);
   const subgraphEndpointMap = new Map<string, string>();
   const subgraphTypesMap = new Map<string, TypeDefinitionNode[]>();
   const typeNameKeysBySubgraphMap = new Map<string, Map<string, string[]>>();
@@ -58,7 +112,7 @@ export function getSubschemasFromSupergraphSdl({
   const orphanTypeMap = new Map<string, TypeDefinitionNode>();
   // TODO: Temporary fix to add missing join__type directives to Query
   const subgraphNames: string[] = [];
-  visit(ast, {
+  visit(supergraphAst, {
     EnumTypeDefinition(node) {
       if (node.name.value === 'join__Graph') {
         node.values?.forEach(valueNode => {
@@ -191,7 +245,7 @@ export function getSubschemasFromSupergraphSdl({
                     extraFields = [];
                     typeNameExtraFieldsMap.set(fieldNodeType.name.value, extraFields);
                   }
-                  const extraFieldTypeNode = ast.definitions.find(
+                  const extraFieldTypeNode = supergraphAst.definitions.find(
                     def => 'name' in def && def.name?.value === fieldNodeType.name.value,
                   ) as ObjectTypeDefinitionNode;
                   providedExtraField.value.value.split(' ').forEach(extraField => {
@@ -311,7 +365,7 @@ export function getSubschemasFromSupergraphSdl({
       orphanTypeMap.set(typeNode.name.value, typeNode);
     }
   }
-  visit(ast, {
+  visit(supergraphAst, {
     ScalarTypeDefinition(node) {
       let isOrphan = !node.name.value.startsWith('link__') && !node.name.value.startsWith('join__');
       node.directives?.forEach(directiveNode => {
@@ -721,6 +775,10 @@ export function getStitchedSchemaFromSupergraphSdl(opts: GetSubschemasFromSuperg
     assumeValidSDL: true,
     typeMergingOptions: {
       useNonNullableFieldOnConflict: true,
+      validationSettings: {
+        validationLevel: ValidationLevel.Off,
+      },
+      fieldConfigMerger: getFieldMergerFromSupergraphSdl(opts.supergraphSdl),
     },
   });
   return filterInternalFieldsAndTypes(supergraphSchema);
diff --git a/packages/stitch/src/createDelegationPlanBuilder.ts b/packages/stitch/src/createDelegationPlanBuilder.ts
@@ -145,17 +145,22 @@ function calculateDelegationStage(
         const fields = typeInSubschema.getFields();
         const field = fields[fieldNode.name.value];
         if (field != null) {
-          const unavailableFields = extractUnavailableFields(field, fieldNode, fieldType => {
-            if (!nonUniqueSubschema.merge?.[fieldType.name]) {
-              delegationMap.set(nonUniqueSubschema, {
-                kind: Kind.SELECTION_SET,
-                selections: [fieldNode],
-              });
-              // Ignore unresolvable fields
-              return false;
-            }
-            return true;
-          });
+          const unavailableFields = extractUnavailableFields(
+            nonUniqueSubschema.transformedSchema,
+            field,
+            fieldNode,
+            fieldType => {
+              if (!nonUniqueSubschema.merge?.[fieldType.name]) {
+                delegationMap.set(nonUniqueSubschema, {
+                  kind: Kind.SELECTION_SET,
+                  selections: [fieldNode],
+                });
+                // Ignore unresolvable fields
+                return false;
+              }
+              return true;
+            },
+          );
           const currentScore = calculateScore(unavailableFields);
           if (currentScore < bestScore) {
             bestScore = currentScore;
@@ -256,7 +261,6 @@ export function createDelegationPlanBuilder(mergedTypeInfo: MergedTypeInfo): Del
       );
       delegationMap = delegationStage.delegationMap;
     }
-
     return delegationMaps;
   });
 }
diff --git a/packages/stitch/src/getFieldsNotInSubschema.ts b/packages/stitch/src/getFieldsNotInSubschema.ts
@@ -4,13 +4,20 @@ import {
   getNamedType,
   GraphQLField,
   GraphQLInterfaceType,
+  GraphQLNamedOutputType,
+  GraphQLNamedType,
   GraphQLObjectType,
   GraphQLSchema,
+  isInterfaceType,
+  isLeafType,
+  isObjectType,
+  isUnionType,
   Kind,
   SelectionNode,
+  SelectionSetNode,
 } from 'graphql';
 import { StitchingInfo } from '@graphql-tools/delegate';
-import { collectSubFields } from '@graphql-tools/utils';
+import { collectSubFields, Maybe } from '@graphql-tools/utils';
 
 export function getFieldsNotInSubschema(
   schema: GraphQLSchema,
@@ -45,6 +52,7 @@ export function getFieldsNotInSubschema(
       const field = fields[fieldName];
       for (const subFieldNode of subFieldNodes) {
         const unavailableFields = extractUnavailableFields(
+          schema,
           field,
           subFieldNode,
           (fieldType, selection) => !fieldNodesByField?.[fieldType.name]?.[selection.name.value],
@@ -77,51 +85,116 @@ export function getFieldsNotInSubschema(
   return Array.from(fieldsNotInSchema);
 }
 
-export function extractUnavailableFields(
-  field: GraphQLField<any, any>,
-  fieldNode: FieldNode,
+export function extractUnavailableFieldsFromSelectionSet(
+  schema: GraphQLSchema,
+  fieldType: GraphQLNamedOutputType,
+  fieldSelectionSet: SelectionSetNode,
   shouldAdd: (fieldType: GraphQLObjectType | GraphQLInterfaceType, selection: FieldNode) => boolean,
 ) {
-  if (fieldNode.selectionSet) {
-    const fieldType = getNamedType(field.type);
-    // TODO: Only object types are supported
-    if (!('getFields' in fieldType)) {
-      return [];
-    }
-    const subFields = fieldType.getFields();
+  if (isLeafType(fieldType)) {
+    return [];
+  }
+  if (isUnionType(fieldType)) {
     const unavailableSelections: SelectionNode[] = [];
-    for (const selection of fieldNode.selectionSet.selections) {
-      if (selection.kind === Kind.FIELD) {
-        if (selection.name.value === '__typename') {
-          continue;
+    for (const type of fieldType.getTypes()) {
+      // Exclude other inline fragments
+      const fieldSelectionExcluded: SelectionSetNode = {
+        ...fieldSelectionSet,
+        selections: fieldSelectionSet.selections.filter(selection =>
+          selection.kind === Kind.INLINE_FRAGMENT
+            ? selection.typeCondition
+              ? selection.typeCondition.name.value === type.name
+              : false
+            : true,
+        ),
+      };
+      unavailableSelections.push(
+        ...extractUnavailableFieldsFromSelectionSet(
+          schema,
+          type,
+          fieldSelectionExcluded,
+          shouldAdd,
+        ),
+      );
+    }
+    return unavailableSelections;
+  }
+  const subFields = fieldType.getFields();
+  const unavailableSelections: SelectionNode[] = [];
+  for (const selection of fieldSelectionSet.selections) {
+    if (selection.kind === Kind.FIELD) {
+      if (selection.name.value === '__typename') {
+        continue;
+      }
+      const fieldName = selection.name.value;
+      const selectionField = subFields[fieldName];
+      if (!selectionField) {
+        if (shouldAdd(fieldType, selection)) {
+          unavailableSelections.push(selection);
         }
-        const fieldName = selection.name.value;
-        const selectionField = subFields[fieldName];
-        if (!selectionField) {
-          if (shouldAdd(fieldType, selection)) {
-            unavailableSelections.push(selection);
-          }
-        } else {
-          const unavailableSubFields = extractUnavailableFields(
-            selectionField,
-            selection,
-            shouldAdd,
-          );
-          if (unavailableSubFields.length) {
-            unavailableSelections.push({
-              ...selection,
-              selectionSet: {
-                kind: Kind.SELECTION_SET,
-                selections: unavailableSubFields,
-              },
-            });
-          }
+      } else {
+        const unavailableSubFields = extractUnavailableFields(
+          schema,
+          selectionField,
+          selection,
+          shouldAdd,
+        );
+        if (unavailableSubFields.length) {
+          unavailableSelections.push({
+            ...selection,
+            selectionSet: {
+              kind: Kind.SELECTION_SET,
+              selections: unavailableSubFields,
+            },
+          });
+        }
+      }
+    } else if (selection.kind === Kind.INLINE_FRAGMENT) {
+      const subFieldType: Maybe<GraphQLNamedType> = selection.typeCondition
+        ? schema.getType(selection.typeCondition.name.value)
+        : fieldType;
+      if (
+        !(isInterfaceType(subFieldType) && isObjectType(subFieldType)) ||
+        subFieldType === fieldType ||
+        (isInterfaceType(fieldType) && schema.isSubType(fieldType, subFieldType))
+      ) {
+        const unavailableFields = extractUnavailableFieldsFromSelectionSet(
+          schema,
+          fieldType,
+          selection.selectionSet,
+          shouldAdd,
+        );
+        if (unavailableFields.length) {
+          unavailableSelections.push({
+            ...selection,
+            selectionSet: {
+              kind: Kind.SELECTION_SET,
+              selections: unavailableFields,
+            },
+          });
         }
-      } else if (selection.kind === Kind.INLINE_FRAGMENT) {
-        // TODO: Support for inline fragments
+      } else {
+        unavailableSelections.push(selection);
       }
     }
-    return unavailableSelections;
+  }
+  return unavailableSelections;
+}
+
+export function extractUnavailableFields(
+  schema: GraphQLSchema,
+  field: GraphQLField<any, any>,
+  fieldNode: FieldNode,
+  shouldAdd: (fieldType: GraphQLObjectType | GraphQLInterfaceType, selection: FieldNode) => boolean,
+) {
+  if (fieldNode.selectionSet) {
+    const fieldType = getNamedType(field.type);
+    return extractUnavailableFieldsFromSelectionSet(
+      schema,
+      fieldType,
+      fieldNode.selectionSet,
+      shouldAdd,
+    );
   }
   return [];
 }
diff --git a/packages/stitch/src/index.ts b/packages/stitch/src/index.ts
@@ -6,3 +6,4 @@ export * from './subschemaConfigTransforms/index.js';
 export * from './types.js';
 export * from './relay.js';
 export * from './executor.js';
+export { getDefaultFieldConfigMerger } from './mergeCandidates.js';
diff --git a/packages/stitch/tests/extractUnavailableFields.test.ts b/packages/stitch/tests/extractUnavailableFields.test.ts
