Add a suspenseCache option to useSuspenseQuery (#10594)

Co-authored-by: Lenz Weber-Tronic <lenz@apollographql.com>
Co-authored-by: Jerel Miller <jerelmiller@gmail.com>

.changeset/curvy-dragons-talk.md

@@ -0,0 +1,5 @@
+---
+'@apollo/client': patch
+---
+
+Add a `suspenseCache` option to `useSuspenseQuery`

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

@@ -298,14 +298,53 @@ describe('useSuspenseQuery', () => {
       });
     }).toThrowError(
       new InvariantError(
-        'Could not find a "suspenseCache" in the context. Wrap the root component ' +
-          'in an <ApolloProvider> and provide a suspenseCache.'
+        'Could not find a "suspenseCache" in the context or passed in as an option. ' +
+          'Wrap the root component in an <ApolloProvider> and provide a suspenseCache, ' +
+          'or pass a SuspenseCache instance in via options.'
       )
     );
 
     consoleSpy.mockRestore();
   });
 
+  it('does not throw when provided a `suspenseCache` option', () => {
+    const { query } = useSimpleQueryCase();
+
+    const client = new ApolloClient({ cache: new InMemoryCache() });
+    const suspenseCache = new SuspenseCache();
+
+    expect(() => {
+      renderHook(() => useSuspenseQuery(query, { suspenseCache }), {
+        wrapper: ({ children }) => (
+          <ApolloProvider client={client} suspenseCache={undefined}>
+            {children}
+          </ApolloProvider>
+        ),
+      });
+    }).not.toThrow();
+  });
+
+  it('prioritizes the `suspenseCache` option over the context value', () => {
+    const { query } = useSimpleQueryCase();
+
+    const directSuspenseCache = new SuspenseCache();
+    const contextSuspenseCache = new SuspenseCache();
+
+    renderHook(
+      () => useSuspenseQuery(query, { suspenseCache: directSuspenseCache }),
+      {
+        wrapper: ({ children }) => (
+          <MockedProvider suspenseCache={contextSuspenseCache}>
+            {children}
+          </MockedProvider>
+        ),
+      }
+    );
+
+    expect(directSuspenseCache.lookup(query, {})).toBeTruthy();
+    expect(contextSuspenseCache.lookup(query, {})).not.toBeTruthy();
+  });
+
   it('ensures a valid fetch policy is used', () => {
     const INVALID_FETCH_POLICIES = ['cache-only', 'standby'];
     const consoleSpy = jest.spyOn(console, 'error').mockImplementation();

src/react/hooks/useSuspenseCache.ts

@@ -1,14 +1,17 @@
 import { useContext } from 'react';
 import { getApolloContext } from '../context';
 import { invariant } from '../../utilities/globals';
+import { SuspenseCache } from '../cache';
 
-export function useSuspenseCache() {
-  const { suspenseCache } = useContext(getApolloContext());
+export function useSuspenseCache(override?: SuspenseCache) {
+  const context = useContext(getApolloContext());
+  const suspenseCache = override || context.suspenseCache;
 
   invariant(
     suspenseCache,
-    'Could not find a "suspenseCache" in the context. Wrap the root component ' +
-      'in an <ApolloProvider> and provide a suspenseCache.'
+    'Could not find a "suspenseCache" in the context or passed in as an option. ' +
+      'Wrap the root component in an <ApolloProvider> and provide a suspenseCache, ' +
+      'or pass a SuspenseCache instance in via options.'
   );
 
   return suspenseCache;

src/react/hooks/useSuspenseQuery.ts

@@ -58,7 +58,7 @@ export function useSuspenseQuery_experimental<
   query: DocumentNode | TypedDocumentNode<TData, TVariables>,
   options: SuspenseQueryHookOptions<TData, TVariables> = Object.create(null)
 ): UseSuspenseQueryResult<TData, TVariables> {
-  const suspenseCache = useSuspenseCache();
+  const suspenseCache = useSuspenseCache(options.suspenseCache);
   const client = useApolloClient(options.client);
   const watchQueryOptions = useWatchQueryOptions({ query, options, client });
   const previousWatchQueryOptionsRef = useRef(watchQueryOptions);

src/react/types/types.ts

@@ -19,6 +19,7 @@ import {
   WatchQueryFetchPolicy,
 } from '../../core';
 import { NextFetchPolicyContext } from '../../core/watchQueryOptions';
+import { SuspenseCache } from '../cache';
 
 /* Common types */
 
@@ -138,6 +139,7 @@ export interface SuspenseQueryHookOptions<
         context: NextFetchPolicyContext<TData, TVariables>
       ) => SuspenseQueryHookFetchPolicy);
   suspensePolicy?: SuspensePolicy;
+  suspenseCache?: SuspenseCache;
 }
 
 /**