Support extensions in useSubscription

This adds optional support for extensions in useSubscription following the graphql-ws spec: https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md
apollographql · May 17, 2024 · 8b9bc7a · 8b9bc7a
1 parent d773000
commit 8b9bc7a
Show file tree

Hide file tree

Showing 8 changed files with 126 additions and 34 deletions.
diff --git a/.api-reports/api-report-core.md b/.api-reports/api-report-core.md
@@ -2124,12 +2124,13 @@ export type SubscribeToMoreOptions<TData = any, TSubscriptionVariables = Operati
 };
 
 // @public (undocumented)
-export interface SubscriptionOptions<TVariables = OperationVariables, TData = any> {
+export interface SubscriptionOptions<TVariables = OperationVariables, TData = any, TExtensions = Record<string, any>> {
     context?: DefaultContext;
     errorPolicy?: ErrorPolicy;
     fetchPolicy?: FetchPolicy;
     query: DocumentNode | TypedDocumentNode<TData, TVariables>;
     variables?: TVariables;
+    extensions?: TExtensions;
 }
 
 // @public (undocumented)

diff --git a/.changeset/angry-seals-jog.md b/.changeset/angry-seals-jog.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client": minor
+---
+
+Support extensions in useSubscription
diff --git a/src/core/QueryManager.ts b/src/core/QueryManager.ts
@@ -293,6 +293,7 @@ export class QueryManager<TStore> {
             optimisticResponse: isOptimistic ? optimisticResponse : void 0,
           },
           variables,
+          {},
           false
         ),
 
@@ -981,52 +982,55 @@ export class QueryManager<TStore> {
     errorPolicy = "none",
     variables,
     context = {},
+    extensions = {},
   }: SubscriptionOptions): Observable<FetchResult<T>> {
     query = this.transform(query);
     variables = this.getVariables(query, variables);
 
     const makeObservable = (variables: OperationVariables) =>
-      this.getObservableFromLink<T>(query, context, variables).map((result) => {
-        if (fetchPolicy !== "no-cache") {
-          // the subscription interface should handle not sending us results we no longer subscribe to.
-          // XXX I don't think we ever send in an object with errors, but we might in the future...
-          if (shouldWriteResult(result, errorPolicy)) {
-            this.cache.write({
-              query,
-              result: result.data,
-              dataId: "ROOT_SUBSCRIPTION",
-              variables: variables,
-            });
+      this.getObservableFromLink<T>(query, context, variables, extensions).map(
+        (result) => {
+          if (fetchPolicy !== "no-cache") {
+            // the subscription interface should handle not sending us results we no longer subscribe to.
+            // XXX I don't think we ever send in an object with errors, but we might in the future...
+            if (shouldWriteResult(result, errorPolicy)) {
+              this.cache.write({
+                query,
+                result: result.data,
+                dataId: "ROOT_SUBSCRIPTION",
+                variables: variables,
+              });
+            }
+
+            this.broadcastQueries();
           }
 
-          this.broadcastQueries();
-        }
+          const hasErrors = graphQLResultHasError(result);
+          const hasProtocolErrors = graphQLResultHasProtocolErrors(result);
+          if (hasErrors || hasProtocolErrors) {
+            const errors: ApolloErrorOptions = {};
+            if (hasErrors) {
+              errors.graphQLErrors = result.errors;
+            }
+            if (hasProtocolErrors) {
+              errors.protocolErrors = result.extensions[PROTOCOL_ERRORS_SYMBOL];
+            }
 
-        const hasErrors = graphQLResultHasError(result);
-        const hasProtocolErrors = graphQLResultHasProtocolErrors(result);
-        if (hasErrors || hasProtocolErrors) {
-          const errors: ApolloErrorOptions = {};
-          if (hasErrors) {
-            errors.graphQLErrors = result.errors;
-          }
-          if (hasProtocolErrors) {
-            errors.protocolErrors = result.extensions[PROTOCOL_ERRORS_SYMBOL];
+            // `errorPolicy` is a mechanism for handling GraphQL errors, according
+            // to our documentation, so we throw protocol errors regardless of the
+            // set error policy.
+            if (errorPolicy === "none" || hasProtocolErrors) {
+              throw new ApolloError(errors);
+            }
           }
 
-          // `errorPolicy` is a mechanism for handling GraphQL errors, according
-          // to our documentation, so we throw protocol errors regardless of the
-          // set error policy.
-          if (errorPolicy === "none" || hasProtocolErrors) {
-            throw new ApolloError(errors);
+          if (errorPolicy === "ignore") {
+            delete result.errors;
           }
-        }
 
-        if (errorPolicy === "ignore") {
-          delete result.errors;
+          return result;
         }
-
-        return result;
-      });
+      );
 
     if (this.getDocumentInfo(query).hasClientExports) {
       const observablePromise = this.localState
@@ -1088,6 +1092,7 @@ export class QueryManager<TStore> {
     query: DocumentNode,
     context: any,
     variables?: OperationVariables,
+    extensions?: Record<string, any>,
     // Prefer context.queryDeduplication if specified.
     deduplication: boolean = context?.queryDeduplication ??
       this.queryDeduplication
@@ -1106,6 +1111,7 @@ export class QueryManager<TStore> {
           ...context,
           forceFetch: !deduplication,
         }),
+        extensions,
       };
 
       context = operation.context;

diff --git a/src/core/watchQueryOptions.ts b/src/core/watchQueryOptions.ts
@@ -206,6 +206,9 @@ export interface SubscriptionOptions<
 
   /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#context:member} */
   context?: DefaultContext;
+
+  /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#extensions:member} */
+  extensions?: Record<string, any>;
 }
 
 export interface MutationBaseOptions<

diff --git a/src/react/hooks/__tests__/useSubscription.test.tsx b/src/react/hooks/__tests__/useSubscription.test.tsx
@@ -454,6 +454,74 @@ describe("useSubscription Hook", () => {
     expect(context!).toBe("Audi");
   });
 
+  it("should share extensions set in options", async () => {
+    const subscription = gql`
+      subscription {
+        car {
+          make
+        }
+      }
+    `;
+
+    const results = ["Audi", "BMW"].map((make) => ({
+      result: { data: { car: { make } } },
+    }));
+
+    let extensions: string;
+    const link = new MockSubscriptionLink();
+    const extensionsLink = new ApolloLink((operation, forward) => {
+      extensions = operation.extensions.make;
+      return forward(operation);
+    });
+    const client = new ApolloClient({
+      link: concat(extensionsLink, link),
+      cache: new Cache({ addTypename: false }),
+    });
+
+    const { result } = renderHook(
+      () =>
+        useSubscription(subscription, {
+          extensions: { make: "Audi" },
+        }),
+      {
+        wrapper: ({ children }) => (
+          <ApolloProvider client={client}>{children}</ApolloProvider>
+        ),
+      }
+    );
+
+    expect(result.current.loading).toBe(true);
+    expect(result.current.error).toBe(undefined);
+    expect(result.current.data).toBe(undefined);
+    setTimeout(() => {
+      link.simulateResult(results[0]);
+    }, 100);
+
+    await waitFor(
+      () => {
+        expect(result.current.data).toEqual(results[0].result.data);
+      },
+      { interval: 1 }
+    );
+    expect(result.current.loading).toBe(false);
+    expect(result.current.error).toBe(undefined);
+
+    setTimeout(() => {
+      link.simulateResult(results[1]);
+    });
+
+    await waitFor(
+      () => {
+        expect(result.current.data).toEqual(results[1].result.data);
+      },
+      { interval: 1 }
+    );
+    expect(result.current.loading).toBe(false);
+    expect(result.current.error).toBe(undefined);
+
+    expect(extensions!).toBe("Audi");
+  });
+
   it("should handle multiple subscriptions properly", async () => {
     const subscription = gql`
       subscription {

diff --git a/src/react/hooks/useSubscription.ts b/src/react/hooks/useSubscription.ts
@@ -146,6 +146,7 @@ export function useSubscription<
       variables: options?.variables,
       fetchPolicy: options?.fetchPolicy,
       context: options?.context,
+      extensions: options?.extensions,
     });
   });
 
@@ -198,6 +199,7 @@ export function useSubscription<
           variables: options?.variables,
           fetchPolicy: options?.fetchPolicy,
           context: options?.context,
+          extensions: options?.extensions,
         })
       );
       canResetObservableRef.current = false;

diff --git a/src/react/types/types.documentation.ts b/src/react/types/types.documentation.ts
@@ -546,6 +546,11 @@ export interface SubscriptionOptionsDocumentation {
    */
   context: unknown;
 
+  /**
+   * Shared context between your component and your network interface (Apollo Link).
+   */
+  extensions: unknown;
+
   /**
    * Allows the registration of a callback function that will be triggered each time the `useSubscription` Hook / `Subscription` component completes the subscription.
    *

diff --git a/src/react/types/types.ts b/src/react/types/types.ts
@@ -447,6 +447,8 @@ export interface BaseSubscriptionOptions<
   skip?: boolean;
   /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#context:member} */
   context?: DefaultContext;
+  /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#extensions:member} */
+  extensions?: Record<string, any>;
   /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#onComplete:member} */
   onComplete?: () => void;
   /** {@inheritDoc @apollo/client!SubscriptionOptionsDocumentation#onData:member} */