Add support for subscribeToMore and client from useSuspenseQuery (

#10450)

* Add support for subscribeToMore from useSuspenseQuery

* Add changeset

* Return the client from `useSuspenseQuery`

* More robust matching on subscribeToMore test

* Update bundlesize

.changeset/wild-mice-nail.md

@@ -0,0 +1,5 @@
+---
+'@apollo/client': patch
+---
+
+Add support for the `subscribeToMore` and `client` fields returned in the `useSuspenseQuery` result.

config/bundlesize.ts

@@ -3,7 +3,7 @@ import { join } from "path";
 import { gzipSync } from "zlib";
 import bytes from "bytes";
 
-const gzipBundleByteLengthLimit = bytes("33.28KB");
+const gzipBundleByteLengthLimit = bytes("33.30KB");
 const minFile = join("dist", "apollo-client.min.cjs");
 const minPath = join(__dirname, "..", minFile);
 const gzipByteLen = gzipSync(readFileSync(minPath)).byteLength;

src/react/hooks/__tests__/useSuspenseQuery.test.tsx

@@ -20,9 +20,16 @@ import {
   DocumentNode,
   InMemoryCache,
   Observable,
+  OperationVariables,
+  SubscribeToMoreOptions,
   TypedDocumentNode,
+  split,
 } from '../../../core';
-import { compact, concatPagination } from '../../../utilities';
+import {
+  compact,
+  concatPagination,
+  getMainDefinition,
+} from '../../../utilities';
 import {
   MockedProvider,
   MockedResponse,
@@ -629,6 +636,28 @@ describe('useSuspenseQuery', () => {
     ]);
   });
 
+  it('returns the client used in the result', async () => {
+    const { query } = useSimpleQueryCase();
+
+    const client = new ApolloClient({
+      link: new ApolloLink(() =>
+        Observable.of({ data: { greeting: 'hello' } })
+      ),
+      cache: new InMemoryCache(),
+    });
+
+    const { result } = renderSuspenseHook(() => useSuspenseQuery(query), {
+      client,
+    });
+
+    // wait for query to finish suspending to avoid warnings
+    await waitFor(() => {
+      expect(result.current.data).toEqual({ greeting: 'hello' });
+    });
+
+    expect(result.current.client).toBe(client);
+  });
+
   it('does not suspend when data is in the cache and using a "cache-first" fetch policy', async () => {
     const { query, mocks } = useSimpleQueryCase();
 
@@ -4906,4 +4935,96 @@ describe('useSuspenseQuery', () => {
       },
     ]);
   });
+
+  it('can subscribe to subscriptions and react to cache updates via `subscribeToMore`', async () => {
+    interface SubscriptionData {
+      greetingUpdated: string;
+    }
+
+    interface QueryData {
+      greeting: string;
+    }
+
+    type UpdateQueryFn = NonNullable<
+      SubscribeToMoreOptions<
+        QueryData,
+        OperationVariables,
+        SubscriptionData
+      >['updateQuery']
+    >;
+
+    const { mocks, query } = useSimpleQueryCase();
+
+    const wsLink = new MockSubscriptionLink();
+    const mockLink = new MockLink(mocks);
+
+    const link = split(
+      ({ query }) => {
+        const definition = getMainDefinition(query);
+
+        return (
+          definition.kind === 'OperationDefinition' &&
+          definition.operation === 'subscription'
+        );
+      },
+      wsLink,
+      mockLink
+    );
+
+    const { result, renders } = renderSuspenseHook(
+      () => useSuspenseQuery(query, { errorPolicy: 'ignore' }),
+      { link }
+    );
+
+    await waitFor(() => {
+      expect(result.current.data).toEqual({ greeting: 'Hello' });
+    });
+
+    const updateQuery = jest.fn<
+      ReturnType<UpdateQueryFn>,
+      Parameters<UpdateQueryFn>
+    >((_, { subscriptionData: { data } }) => {
+      return { greeting: data.greetingUpdated };
+    });
+
+    result.current.subscribeToMore<SubscriptionData>({
+      document: gql`
+        subscription {
+          greetingUpdated
+        }
+      `,
+      updateQuery,
+    });
+
+    wsLink.simulateResult({
+      result: {
+        data: {
+          greetingUpdated: 'Subscription hello',
+        },
+      },
+    });
+
+    await waitFor(() => {
+      expect(result.current.data).toEqual({
+        greeting: 'Subscription hello',
+      });
+    });
+
+    expect(updateQuery).toHaveBeenCalledTimes(1);
+    expect(updateQuery).toHaveBeenCalledWith(
+      { greeting: 'Hello' },
+      {
+        subscriptionData: {
+          data: { greetingUpdated: 'Subscription hello' },
+        },
+        variables: {},
+      }
+    );
+
+    expect(renders.count).toBe(3);
+    expect(renders.frames).toMatchObject([
+      { data: { greeting: 'Hello' } },
+      { data: { greeting: 'Subscription hello' } },
+    ]);
+  });
 });

src/react/hooks/useSuspenseQuery.ts

@@ -1,10 +1,4 @@
-import {
-  useRef,
-  useEffect,
-  useCallback,
-  useMemo,
-  useState,
-} from 'react';
+import { useRef, useEffect, useCallback, useMemo, useState } from 'react';
 import { equal } from '@wry/equality';
 import {
   ApolloClient,
@@ -38,10 +32,12 @@ export interface UseSuspenseQueryResult<
   TData = any,
   TVariables = OperationVariables
 > {
+  client: ApolloClient<any>;
   data: TData;
   error: ApolloError | undefined;
   fetchMore: ObservableQueryFields<TData, TVariables>['fetchMore'];
   refetch: ObservableQueryFields<TData, TVariables>['refetch'];
+  subscribeToMore: ObservableQueryFields<TData, TVariables>['subscribeToMore'];
 }
 
 const SUPPORTED_FETCH_POLICIES: WatchQueryFetchPolicy[] = [
@@ -151,6 +147,7 @@ export function useSuspenseQuery_experimental<
 
   return useMemo(() => {
     return {
+      client,
       data: result.data,
       error: errorPolicy === 'ignore' ? void 0 : toApolloError(result),
       fetchMore: (options) => {
@@ -173,8 +170,9 @@ export function useSuspenseQuery_experimental<
 
         return promise;
       },
+      subscribeToMore: (options) => observable.subscribeToMore(options),
     };
-  }, [result, observable, errorPolicy]);
+  }, [client, result, observable, errorPolicy]);
 }
 
 function validateOptions(options: WatchQueryOptions) {