enhance(delegate): improve request delegation (#6576)

* enhance(delegate): improve request delegation * More * Fix typechecking
ardatan · Oct 16, 2024 · 4cdb462 · 4cdb462
1 parent 7afe685
commit 4cdb462
Show file tree

Hide file tree

Showing 5 changed files with 227 additions and 176 deletions.
diff --git a/.changeset/kind-dolls-wonder.md b/.changeset/kind-dolls-wonder.md
@@ -0,0 +1,5 @@
+---
+'@graphql-tools/delegate': patch
+---
+
+Performance improvements on upstream request execution
diff --git a/packages/delegate/src/OverlappingAliasesTransform.ts b/packages/delegate/src/OverlappingAliasesTransform.ts
@@ -1,5 +1,5 @@
-import { isNullableType, Kind, visit } from 'graphql';
-import { ExecutionRequest, ExecutionResult } from '@graphql-tools/utils';
+import { ASTNode, isNullableType, Kind, visit } from 'graphql';
+import { ASTVisitorKeyMap, ExecutionRequest, ExecutionResult } from '@graphql-tools/utils';
 import { DelegationContext, Transform } from './types.js';
 
 const OverlappingAliases = Symbol('OverlappingAliases');
@@ -16,65 +16,80 @@ export class OverlappingAliasesTransform<TContext>
     delegationContext: DelegationContext<TContext>,
     transformationContext: OverlappingAliasesContext,
   ) {
-    const newDocument = visit(request.document, {
-      [Kind.SELECTION_SET]: node => {
-        const seenNonNullable = new Set<string>();
-        const seenNullable = new Set<string>();
-        return {
-          ...node,
-          selections: node.selections.map(selection => {
-            if (selection.kind === Kind.INLINE_FRAGMENT) {
-              const selectionTypeName = selection.typeCondition?.name.value;
-              if (selectionTypeName) {
-                const selectionType =
-                  delegationContext.transformedSchema.getType(selectionTypeName);
-                if (selectionType && 'getFields' in selectionType) {
-                  const selectionTypeFields = selectionType.getFields();
-                  return {
-                    ...selection,
-                    selectionSet: {
-                      ...selection.selectionSet,
-                      selections: selection.selectionSet.selections.map(subSelection => {
-                        if (subSelection.kind === Kind.FIELD) {
-                          const fieldName = subSelection.name.value;
-                          if (!subSelection.alias) {
-                            const field = selectionTypeFields[fieldName];
-                            if (field) {
-                              let currentNullable: boolean;
-                              if (isNullableType(field.type)) {
-                                seenNullable.add(fieldName);
-                                currentNullable = true;
-                              } else {
-                                seenNonNullable.add(fieldName);
-                                currentNullable = false;
-                              }
-                              if (seenNullable.has(fieldName) && seenNonNullable.has(fieldName)) {
-                                transformationContext[OverlappingAliases] = true;
-                                return {
-                                  ...subSelection,
-                                  alias: {
-                                    kind: Kind.NAME,
-                                    value: currentNullable
-                                      ? `_nullable_${fieldName}`
-                                      : `_nonNullable_${fieldName}`,
-                                  },
-                                };
-                              }
+    const visitorKeys: ASTVisitorKeyMap = {
+      Document: ['definitions'],
+      OperationDefinition: ['selectionSet'],
+      SelectionSet: ['selections'],
+      Field: ['selectionSet'],
+      InlineFragment: ['selectionSet'],
+      FragmentDefinition: ['selectionSet'],
+    };
+    const seenNonNullableMap = new WeakMap<readonly ASTNode[], Set<string>>();
+    const seenNullableMap = new WeakMap<readonly ASTNode[], Set<string>>();
+    const newDocument = visit(
+      request.document,
+      {
+        [Kind.INLINE_FRAGMENT](selection, _key, parent, _path, _ancestors) {
+          if (Array.isArray(parent)) {
+            const selectionTypeName = selection.typeCondition?.name.value;
+            if (selectionTypeName) {
+              const selectionType = delegationContext.transformedSchema.getType(selectionTypeName);
+              if (selectionType && 'getFields' in selectionType) {
+                const selectionTypeFields = selectionType.getFields();
+                let seenNonNullable = seenNonNullableMap.get(parent);
+                if (!seenNonNullable) {
+                  seenNonNullable = new Set();
+                  seenNonNullableMap.set(parent, seenNonNullable);
+                }
+                let seenNullable = seenNullableMap.get(parent);
+                if (!seenNullable) {
+                  seenNullable = new Set();
+                  seenNullableMap.set(parent, seenNullable);
+                }
+                return {
+                  ...selection,
+                  selectionSet: {
+                    ...selection.selectionSet,
+                    selections: selection.selectionSet.selections.map(subSelection => {
+                      if (subSelection.kind === Kind.FIELD) {
+                        const fieldName = subSelection.name.value;
+                        if (!subSelection.alias) {
+                          const field = selectionTypeFields[fieldName];
+                          if (field) {
+                            let currentNullable: boolean;
+                            if (isNullableType(field.type)) {
+                              seenNullable.add(fieldName);
+                              currentNullable = true;
+                            } else {
+                              seenNonNullable.add(fieldName);
+                              currentNullable = false;
+                            }
+                            if (seenNullable.has(fieldName) && seenNonNullable.has(fieldName)) {
+                              transformationContext[OverlappingAliases] = true;
+                              return {
+                                ...subSelection,
+                                alias: {
+                                  kind: Kind.NAME,
+                                  value: currentNullable
+                                    ? `_nullable_${fieldName}`
+                                    : `_nonNullable_${fieldName}`,
+                                },
+                              };
                             }
                           }
                         }
-                        return subSelection;
-                      }),
-                    },
-                  };
-                }
+                      }
+                      return subSelection;
+                    }),
+                  },
+                };
               }
             }
-            return selection;
-          }),
-        };
+          }
+        },
       },
-    });
+      visitorKeys as any,
+    );
     return {
       ...request,
       document: newDocument,

diff --git a/packages/delegate/src/finalizeGatewayRequest.ts b/packages/delegate/src/finalizeGatewayRequest.ts
@@ -174,44 +174,58 @@ export function finalizeGatewayRequest<TContext>(
   let cleanedUpDocument = newDocument;
 
   // TODO: Optimize this internally later
-  cleanedUpDocument = visit(newDocument, {
-    // Cleanup extra __typename fields
-    SelectionSet: {
-      leave(node) {
-        const { hasTypeNameField, selections } = filterTypenameFields(node.selections);
-        if (hasTypeNameField) {
-          selections.unshift({
-            kind: Kind.FIELD,
-            name: {
-              kind: Kind.NAME,
-              value: '__typename',
-            },
-          });
-        }
-        return {
-          ...node,
-          selections,
-        };
-      },
-    },
-    // Cleanup empty inline fragments
-    InlineFragment: {
-      leave(node) {
-        // No need __typename in inline fragment
-        const { selections } = filterTypenameFields(node.selectionSet.selections);
-        if (selections.length === 0) {
-          return null;
-        }
-        return {
-          ...node,
-          selectionSet: {
-            ...node.selectionSet,
+  const visitorKeys: ASTVisitorKeyMap = {
+    Document: ['definitions'],
+    OperationDefinition: ['selectionSet'],
+    SelectionSet: ['selections'],
+    Field: ['selectionSet'],
+    InlineFragment: ['selectionSet'],
+    FragmentDefinition: ['selectionSet'],
+  };
+  cleanedUpDocument = visit(
+    newDocument,
+    {
+      // Cleanup extra __typename fields
+      [Kind.SELECTION_SET]: {
+        leave(node) {
+          const { hasTypeNameField, selections } = filterTypenameFields(node.selections);
+          if (hasTypeNameField) {
+            selections.unshift({
+              kind: Kind.FIELD,
+              name: {
+                kind: Kind.NAME,
+                value: '__typename',
+              },
+            });
+          }
+          return {
+            ...node,
             selections,
-          },
-        };
+          };
+        },
+      },
+      // Cleanup empty inline fragments
+      [Kind.INLINE_FRAGMENT]: {
+        leave(node) {
+          // No need __typename in inline fragment
+          const { selections } = filterTypenameFields(node.selectionSet.selections);
+          if (selections.length === 0) {
+            return null;
+          }
+          return {
+            ...node,
+            selectionSet: {
+              ...node.selectionSet,
+              selections,
+            },
+            // @defer is not available for the communication between the gw and subgraph
+            directives: node.directives?.filter?.(directive => directive.name.value !== 'defer'),
+          };
+        },
       },
     },
-  });
+    visitorKeys as any,
+  );
 
   return {
     ...originalRequest,

diff --git a/packages/delegate/src/prepareGatewayDocument.ts b/packages/delegate/src/prepareGatewayDocument.ts
@@ -46,22 +46,31 @@ export function prepareGatewayDocument(
   }
 
   const visitedSelections = new WeakSet<SelectionNode>();
-  wrappedConcreteTypesDocument = visit(wrappedConcreteTypesDocument, {
-    [Kind.SELECTION_SET](node) {
-      const newSelections: Array<SelectionNode> = [];
-      for (const selectionNode of node.selections) {
-        if (
-          selectionNode.kind === Kind.INLINE_FRAGMENT &&
-          selectionNode.typeCondition != null &&
-          !visitedSelections.has(selectionNode)
-        ) {
-          visitedSelections.add(selectionNode);
-          const typeName = selectionNode.typeCondition.name.value;
-          const gatewayType = infoSchema.getType(typeName);
-          const subschemaType = transformedSchema.getType(typeName);
-          if (isAbstractType(gatewayType)) {
-            const possibleTypes = infoSchema.getPossibleTypes(gatewayType);
-            if (isAbstractType(subschemaType)) {
+  const visitorKeys: ASTVisitorKeyMap = {
+    Document: ['definitions'],
+    OperationDefinition: ['selectionSet'],
+    SelectionSet: ['selections'],
+    Field: ['selectionSet'],
+    InlineFragment: ['selectionSet'],
+    FragmentDefinition: ['selectionSet'],
+  };
+  wrappedConcreteTypesDocument = visit(
+    wrappedConcreteTypesDocument,
+    {
+      [Kind.SELECTION_SET](node) {
+        const newSelections: Array<SelectionNode> = [];
+        for (const selectionNode of node.selections) {
+          if (
+            selectionNode.kind === Kind.INLINE_FRAGMENT &&
+            selectionNode.typeCondition != null &&
+            !visitedSelections.has(selectionNode)
+          ) {
+            visitedSelections.add(selectionNode);
+            const typeName = selectionNode.typeCondition.name.value;
+            const gatewayType = infoSchema.getType(typeName);
+            const subschemaType = transformedSchema.getType(typeName);
+            if (isAbstractType(gatewayType) && isAbstractType(subschemaType)) {
+              const possibleTypes = infoSchema.getPossibleTypes(gatewayType);
               const possibleTypesInSubschema = transformedSchema.getPossibleTypes(subschemaType);
               const extraTypesForSubschema = new Set<string>();
               for (const possibleType of possibleTypes) {
@@ -93,34 +102,35 @@ export function prepareGatewayDocument(
                 });
               }
             }
-          }
-          const typeInSubschema = transformedSchema.getType(typeName);
-          if (!typeInSubschema) {
-            for (const selection of selectionNode.selectionSet.selections) {
-              newSelections.push(selection);
+            const typeInSubschema = transformedSchema.getType(typeName);
+            if (!typeInSubschema) {
+              for (const selection of selectionNode.selectionSet.selections) {
+                newSelections.push(selection);
+              }
             }
-          }
-          if (typeInSubschema && 'getFields' in typeInSubschema) {
-            const fieldMap = typeInSubschema.getFields();
-            for (const selection of selectionNode.selectionSet.selections) {
-              if (selection.kind === Kind.FIELD) {
-                const fieldName = selection.name.value;
-                const field = fieldMap[fieldName];
-                if (!field) {
-                  newSelections.push(selection);
+            if (typeInSubschema && 'getFields' in typeInSubschema) {
+              const fieldMap = typeInSubschema.getFields();
+              for (const selection of selectionNode.selectionSet.selections) {
+                if (selection.kind === Kind.FIELD) {
+                  const fieldName = selection.name.value;
+                  const field = fieldMap[fieldName];
+                  if (!field) {
+                    newSelections.push(selection);
+                  }
                 }
               }
             }
           }
+          newSelections.push(selectionNode);
         }
-        newSelections.push(selectionNode);
-      }
-      return {
-        ...node,
-        selections: newSelections,
-      };
+        return {
+          ...node,
+          selections: newSelections,
+        };
+      },
     },
-  });
+    visitorKeys as any,
+  );
 
   const {
     possibleTypesMap,