Fix wrapping subscription types (#5572)

ardatan · Sep 7, 2023 · aadb591 · aadb591
1 parent b8b8164
commit aadb591
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 12 deletions.
diff --git a/.changeset/nine-students-approve.md b/.changeset/nine-students-approve.md
@@ -0,0 +1,6 @@
+---
+'@graphql-tools/delegate': patch
+'@graphql-tools/wrap': patch
+---
+
+Fix for wrapping subscription types
diff --git a/packages/delegate/src/finalizeGatewayRequest.ts b/packages/delegate/src/finalizeGatewayRequest.ts
@@ -89,6 +89,14 @@ function finalizeGatewayDocument(
         operationOrFragmentVariables.indexOf(variable.variable.name.value) !== -1,
     );
 
+    // Prevent unnecessary __typename in Subscription
+    if (operation.operation === 'subscription') {
+      selectionSet.selections = selectionSet.selections.filter(
+        (selection: SelectionNode) =>
+          selection.kind !== Kind.FIELD || selection.name.value !== '__typename',
+      );
+    }
+
     newOperations.push({
       kind: Kind.OPERATION_DEFINITION,
       operation: operation.operation,

diff --git a/packages/stitch/tests/alternateStitchSchemas.test.ts b/packages/stitch/tests/alternateStitchSchemas.test.ts
@@ -38,6 +38,8 @@ import {
   wrapSchema,
   WrapType,
 } from '@graphql-tools/wrap';
+import { Repeater } from '@repeaterjs/repeater';
+import { assertAsyncIterable } from '../../loaders/url/tests/test-utils.js';
 import {
   bookingSchema,
   propertySchema,
@@ -1051,6 +1053,62 @@ describe('WrapType', () => {
     expect(result).toEqual(expectedResult);
   });
 
+  test('Subscription transform should work', async () => {
+    const transformedSchema = wrapSchema({
+      schema: makeExecutableSchema({
+        typeDefs: /* GraphQL */ `
+          type Query {
+            _: Boolean
+          }
+          type Subscription {
+            countdown(start: Int!): Int!
+          }
+        `,
+        resolvers: {
+          Subscription: {
+            countdown: {
+              subscribe: (_root, args) =>
+                new Repeater((push, stop) => {
+                  let counter = args.start;
+                  const interval = setInterval(() => {
+                    push(counter--);
+                    if (counter < 0) {
+                      stop();
+                    }
+                  }, 1000);
+                  return stop.then(() => clearInterval(interval));
+                }),
+              resolve: cnt => cnt,
+            },
+          },
+        },
+      }),
+      transforms: [new WrapType('Subscription', 'Namespace_Subscription', 'namespace')],
+    });
+
+    const result = await subscribe({
+      schema: transformedSchema,
+      document: parse(/* GraphQL */ `
+        subscription {
+          namespace {
+            countdown(start: 3)
+          }
+        }
+      `),
+    });
+    assertAsyncIterable(result);
+    const values = [];
+    for await (const value of result) {
+      values.push(value);
+    }
+    expect(values).toEqual([
+      { data: { namespace: { countdown: 3 } } },
+      { data: { namespace: { countdown: 2 } } },
+      { data: { namespace: { countdown: 1 } } },
+      { data: { namespace: { countdown: 0 } } },
+    ]);
+  });
+
   test('namespacing different subschemas with overlapping root field names', async () => {
     const typeDefGen = (i: number) => `
       type Query {

diff --git a/packages/wrap/src/transforms/WrapFields.ts b/packages/wrap/src/transforms/WrapFields.ts
@@ -131,35 +131,54 @@ export default class WrapFields<TContext extends Record<string, any>>
       wrappingFieldName = this.wrappingFieldNames[wrapIndex];
     }
 
-    const wrappingRootField =
-      this.outerTypeName === originalWrappingSchema.getQueryType()?.name ||
-      this.outerTypeName === originalWrappingSchema.getMutationType()?.name;
+    const targetSchema = subschemaConfig.schema;
+    let wrappingOperation: 'query' | 'mutation' | 'subscription' | undefined;
+    switch (this.outerTypeName) {
+      case targetSchema.getQueryType()?.name:
+        wrappingOperation = 'query';
+        break;
+      case targetSchema.getMutationType()?.name:
+        wrappingOperation = 'mutation';
+        break;
+      case targetSchema.getSubscriptionType()?.name:
+        wrappingOperation = 'subscription';
+        break;
+    }
 
     let resolve: GraphQLFieldResolver<any, any> | undefined;
-    if (wrappingRootField) {
-      const targetSchema = subschemaConfig.schema;
-      const operation =
-        this.outerTypeName === targetSchema.getQueryType()?.name ? 'query' : 'mutation';
+    if (wrappingOperation) {
       const createProxyingResolver =
         subschemaConfig.createProxyingResolver ?? defaultCreateProxyingResolver;
       resolve = createProxyingResolver({
         subschemaConfig,
-        operation: operation as OperationTypeNode,
+        operation: wrappingOperation as OperationTypeNode,
         fieldName: wrappingFieldName,
       });
     } else {
       resolve = defaultMergedResolver;
     }
 
+    const wrappingType = new GraphQLNonNull(
+      newSchema.getType(wrappingTypeName) as GraphQLObjectType,
+    );
+    const newFieldConfig: GraphQLFieldConfig<any, any> =
+      wrappingOperation === 'subscription'
+        ? {
+            type: wrappingType,
+            subscribe: resolve,
+            resolve: (payload: any) => payload,
+          }
+        : {
+            type: wrappingType,
+            resolve,
+          };
+
     [newSchema] = modifyObjectFields(
       newSchema,
       this.outerTypeName,
       fieldName => !!newTargetFieldConfigMap[fieldName],
       {
-        [wrappingFieldName]: {
-          type: new GraphQLNonNull(newSchema.getType(wrappingTypeName) as GraphQLObjectType),
-          resolve,
-        },
+        [wrappingFieldName]: newFieldConfig,
       },
     );