From b1a2f8490889e61c3469a7f2250c2c22b053fd26 Mon Sep 17 00:00:00 2001
From: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
Date: Wed, 6 Mar 2024 14:52:16 +0100
Subject: [PATCH] allow enabling "forced suspense" for `useQuery` hook in SSR

---
 .../dynamic/useQuery_forcedSuspense/page.tsx  | 40 ++++++++++++++++++
 .../client-react-streaming/package-shape.json |  3 ++
 .../WrappedApolloClient.tsx                   |  4 +-
 .../src/DataTransportAbstraction/hooks.ts     | 42 +++++++++++++++++++
 .../src/DataTransportAbstraction/index.ts     |  1 +
 packages/client-react-streaming/src/index.ts  |  1 +
 6 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 integration-test/nextjs/src/app/cc/dynamic/useQuery_forcedSuspense/page.tsx

diff --git a/integration-test/nextjs/src/app/cc/dynamic/useQuery_forcedSuspense/page.tsx b/integration-test/nextjs/src/app/cc/dynamic/useQuery_forcedSuspense/page.tsx
new file mode 100644
index 00000000..2a40a9c0
--- /dev/null
+++ b/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>
+  );
+}
diff --git a/packages/client-react-streaming/package-shape.json b/packages/client-react-streaming/package-shape.json
index 458be7b3..177635d1 100644
--- a/packages/client-react-streaming/package-shape.json
+++ b/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"
     ]
   },
diff --git a/packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.tsx b/packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.tsx
index b8198892..087296df 100644
--- a/packages/client-react-streaming/src/DataTransportAbstraction/WrappedApolloClient.tsx
+++ b/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.
diff --git a/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts b/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts
index c08b0fdc..5e57c2e7 100644
--- a/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts
+++ b/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"]
+          );
+      }
+    : () => {};
diff --git a/packages/client-react-streaming/src/DataTransportAbstraction/index.ts b/packages/client-react-streaming/src/DataTransportAbstraction/index.ts
index 1c16db28..684e5763 100644
--- a/packages/client-react-streaming/src/DataTransportAbstraction/index.ts
+++ b/packages/client-react-streaming/src/DataTransportAbstraction/index.ts
@@ -12,3 +12,4 @@ export {
   WrapApolloProvider,
   WrappedApolloProvider,
 } from "./WrapApolloProvider.js";
+export { enableSSRWaitForUseQuery } from "./hooks.js";
diff --git a/packages/client-react-streaming/src/index.ts b/packages/client-react-streaming/src/index.ts
index 22c3c327..4092dafb 100644
--- a/packages/client-react-streaming/src/index.ts
+++ b/packages/client-react-streaming/src/index.ts
@@ -6,4 +6,5 @@ export {
   QueryEvent,
   WrapApolloProvider,
   WrappedApolloProvider,
+  enableSSRWaitForUseQuery,
 } from "./DataTransportAbstraction/index.js";