allow enabling "forced suspense" for useQuery hook in SSR

Author: phryneas

integration-test/nextjs/src/app/cc/dynamic/useQuery_forcedSuspense/page.tsx

@@ -0,0 +1,40 @@
+"use client";
+
+import { useQuery } from "@apollo/experimental-nextjs-app-support/ssr";
+import { enableSSRWaitForUseQuery } from "@apollo/client-react-streaming";
+import type { TypedDocumentNode } from "@apollo/client";
+import { gql, useApolloClient } from "@apollo/client";
+
+const QUERY: TypedDocumentNode<{
+  products: {
+    id: string;
+    title: string;
+  }[];
+}> = gql`
+  query dynamicProducts {
+    products {
+      id
+      title
+    }
+  }
+`;
+
+export const dynamic = "force-dynamic";
+
+export default function Page() {
+  enableSSRWaitForUseQuery(useApolloClient());
+  const result = useQuery(QUERY);
+  globalThis.hydrationFinished?.();
+
+  if (!result.data) {
+    return <div>Loading...</div>;
+  }
+
+  return (
+    <ul>
+      {result.data.products.map(({ id, title }) => (
+        <li key={id}>{title}</li>
+      ))}
+    </ul>
+  );
+}

packages/client-react-streaming/package-shape.json

@@ -18,6 +18,7 @@
       "SSRMultipartLink",
       "WrapApolloProvider",
       "resetApolloSingletons",
+      "enableSSRWaitForUseQuery",
       "built_for_browser"
     ],
     "node": [
@@ -29,6 +30,7 @@
       "SSRMultipartLink",
       "WrapApolloProvider",
       "resetApolloSingletons",
+      "enableSSRWaitForUseQuery",
       "built_for_ssr"
     ],
     "edge-light,worker,browser": [
@@ -40,6 +42,7 @@
       "SSRMultipartLink",
       "WrapApolloProvider",
       "resetApolloSingletons",
+      "enableSSRWaitForUseQuery",
       "built_for_ssr"
     ]
   },

packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.tsx

@@ -25,7 +25,7 @@ import type {
 } from "./DataTransportAbstraction.js";
 import { bundle } from "../bundleInfo.js";
 
-function getQueryManager<TCacheShape>(
+export function getQueryManager<TCacheShape>(
   client: OrigApolloClient<unknown>
 ): QueryManager<TCacheShape> & {
   [wrappers]: HookWrappers;
@@ -39,7 +39,7 @@ type SimulatedQueryInfo = {
   options: WatchQueryOptions<OperationVariables, any>;
 };
 
-const wrappers = Symbol.for("apollo.hook.wrappers");
+export const wrappers = Symbol.for("apollo.hook.wrappers");
 class ApolloClientBase<TCacheShape> extends OrigApolloClient<TCacheShape> {
   /**
    * Information about the current package and it's export names, for use in error messages.

packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts

@@ -1,5 +1,12 @@
 import type { HookWrappers } from "@apollo/client/react/internal/index.js";
 import { useTransportValue } from "./useTransportValue.js";
+import type { WatchQueryOptions } from "@apollo/client/index.js";
+import { useApolloClient } from "@apollo/client/index.js";
+import { getSuspenseCache } from "@apollo/client/react/internal/index.js";
+import { canonicalStringify } from "@apollo/client/cache/index.js";
+import { use } from "react";
+import type { ApolloClient } from "./WrappedApolloClient.js";
+import { getQueryManager, wrappers } from "./WrappedApolloClient.js";
 
 export const hookWrappers: HookWrappers = {
   useFragment(orig_useFragment) {
@@ -35,3 +42,38 @@ function wrap<T extends (...args: any[]) => any>(
     return { ...result, ...useTransportValue(transported) };
   }) as T;
 }
+
+export const enableSSRWaitForUseQuery: (client: ApolloClient<any>) => void =
+  process.env.REACT_ENV === "ssr"
+    ? (client) => {
+        getQueryManager(client)[wrappers].useQuery = (orig_useQuery) =>
+          wrap<typeof orig_useQuery>(
+            function useQuery(query, options) {
+              const client = useApolloClient();
+              const result = client.cache.read({
+                query,
+                variables: options?.variables,
+                returnPartialData: options?.returnPartialData,
+                optimistic: false,
+              });
+              if (!result) {
+                const queryRef = getSuspenseCache(client).getQueryRef(
+                  [query, canonicalStringify(options?.variables), "useQuery"],
+                  () =>
+                    client.watchQuery({
+                      query,
+                      ...(options as Partial<WatchQueryOptions>),
+                    })
+                );
+                use(queryRef.promise);
+              }
+
+              return orig_useQuery(query, {
+                ...options,
+                fetchPolicy: "cache-only",
+              });
+            },
+            ["data", "loading", "networkStatus", "called"]
+          );
+      }
+    : () => {};

packages/client-react-streaming/src/DataTransportAbstraction/index.ts

@@ -12,3 +12,4 @@ export {
   WrapApolloProvider,
   WrappedApolloProvider,
 } from "./WrapApolloProvider.js";
+export { enableSSRWaitForUseQuery } from "./hooks.js";

packages/client-react-streaming/src/index.ts

@@ -6,4 +6,5 @@ export {
   QueryEvent,
   WrapApolloProvider,
   WrappedApolloProvider,
+  enableSSRWaitForUseQuery,
 } from "./DataTransportAbstraction/index.js";