From 563db9bba32c0b656da0465f60e61abde23af445 Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Tue, 11 Feb 2025 03:08:19 -0700
Subject: [PATCH] Add support for `useSuspenseFragment` (#427)

* Add useSuspenseFragment to hooks

* Update nextjs preload query w/ defer to use useSuspenseFragment

* add changeset

---------

Co-authored-by: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
---
 .changeset/rich-seahorses-reply.md            |  5 ++
 .../ClientChild.tsx                           | 60 ++++++++++++-------
 .../defer-queryRef-useReadQuery/page.tsx      |  4 +-
 integration-test/shared/queries.ts            | 34 ++++++++++-
 .../src/DataTransportAbstraction/hooks.ts     |  3 +
 5 files changed, 80 insertions(+), 26 deletions(-)
 create mode 100644 .changeset/rich-seahorses-reply.md

diff --git a/.changeset/rich-seahorses-reply.md b/.changeset/rich-seahorses-reply.md
new file mode 100644
index 00000000..a8d7e706
--- /dev/null
+++ b/.changeset/rich-seahorses-reply.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client-react-streaming": patch
+---
+
+add support for `useSuspenseFragment`
diff --git a/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/ClientChild.tsx b/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/ClientChild.tsx
index ebde96e5..d811d392 100644
--- a/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/ClientChild.tsx
+++ b/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/ClientChild.tsx
@@ -4,10 +4,14 @@ import {
   useApolloClient,
   useQueryRefHandlers,
   useReadQuery,
+  useSuspenseFragment,
 } from "@apollo/client";
-import { DeferredDynamicProductResult } from "@integration-test/shared/queries";
+import {
+  DeferredDynamicProductResult,
+  RATING_FRAGMENT,
+} from "@integration-test/shared/queries";
 import { TransportedQueryRef } from "@apollo/experimental-nextjs-app-support";
-import { useTransition } from "react";
+import { Suspense, useTransition } from "react";
 
 export function ClientChild({
   queryRef,
@@ -17,21 +21,18 @@ export function ClientChild({
   const { refetch } = useQueryRefHandlers(queryRef);
   const [refetching, startTransition] = useTransition();
   const { data } = useReadQuery(queryRef);
-  const client = useApolloClient();
 
   return (
     <>
       <ul>
-        {data.products.map(({ id, title, rating }) => (
-          <li key={id}>
-            {title}
+        {data.products.map((product) => (
+          <li key={product.id}>
+            {product.title}
             <br />
             Rating:{" "}
-            <div style={{ display: "inline-block", verticalAlign: "text-top" }}>
-              {rating?.value || ""}
-              <br />
-              {rating ? `Queried in ${rating.env} environment` : "loading..."}
-            </div>
+            <Suspense fallback="Loading rating...">
+              <Rating product={product} delayDeferred={1000} />
+            </Suspense>
           </li>
         ))}
       </ul>
@@ -39,18 +40,6 @@ export function ClientChild({
       <button
         disabled={refetching}
         onClick={() => {
-          client.cache.batch({
-            update(cache) {
-              for (const product of data.products) {
-                cache.modify({
-                  id: cache.identify(product),
-                  fields: {
-                    rating: () => null,
-                  },
-                });
-              }
-            },
-          });
           startTransition(() => {
             refetch();
           });
@@ -61,3 +50,28 @@ export function ClientChild({
     </>
   );
 }
+
+interface RatingProps {
+  product: { __typename: "Product"; id: string };
+  delayDeferred: number;
+}
+
+function Rating({ product, delayDeferred }: RatingProps) {
+  const { data } = useSuspenseFragment({
+    fragment: RATING_FRAGMENT,
+    from: product,
+    variables: {
+      delayDeferred,
+    },
+  });
+
+  const { rating } = data;
+
+  return (
+    <div style={{ display: "inline-block", verticalAlign: "text-top" }}>
+      {rating.value || ""}
+      <br />
+      Queried in {rating.env} environment
+    </div>
+  );
+}
diff --git a/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/page.tsx b/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/page.tsx
index 3d538bdf..850fccde 100644
--- a/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/page.tsx
+++ b/integration-test/nextjs/src/app/rsc/dynamic/PreloadQuery/defer-queryRef-useReadQuery/page.tsx
@@ -1,6 +1,6 @@
 import { ApolloWrapper } from "@/app/cc/ApolloWrapper";
 import { ClientChild } from "./ClientChild";
-import { DEFERRED_QUERY } from "@integration-test/shared/queries";
+import { DEFERRED_QUERY_WITH_FRAGMENT } from "@integration-test/shared/queries";
 
 export const dynamic = "force-dynamic";
 import { PreloadQuery } from "../../../client";
@@ -10,7 +10,7 @@ export default async function Page({ searchParams }: { searchParams?: any }) {
   return (
     <ApolloWrapper>
       <PreloadQuery
-        query={DEFERRED_QUERY}
+        query={DEFERRED_QUERY_WITH_FRAGMENT}
         context={{
           delay: 1000,
           error: (await searchParams)?.errorIn || undefined,
diff --git a/integration-test/shared/queries.ts b/integration-test/shared/queries.ts
index e3d73845..66ba3bc9 100644
--- a/integration-test/shared/queries.ts
+++ b/integration-test/shared/queries.ts
@@ -22,9 +22,10 @@ export const QUERY: TypedDocumentNode<
 
 export interface DeferredDynamicProductResult {
   products: {
+    __typename: "Product";
     id: string;
     title: string;
-    rating: undefined | { value: string; env: string };
+    rating: undefined | { __typename: "Rating"; value: string; env: string };
   }[];
   env: string;
 }
@@ -46,3 +47,34 @@ export const DEFERRED_QUERY: TypedDocumentNode<
     env
   }
 `;
+
+export const RATING_FRAGMENT: TypedDocumentNode<
+  {
+    __typename: "Product";
+    rating: { __typename: "Rating"; value: string; env: string };
+  },
+  { delayDeferred: number }
+> = gql`
+  fragment RatingFragment on Product {
+    rating(delay: $delayDeferred) {
+      value
+      env
+    }
+  }
+`;
+
+export const DEFERRED_QUERY_WITH_FRAGMENT: TypedDocumentNode<
+  DeferredDynamicProductResult,
+  { someArgument?: string; delayDeferred: number }
+> = gql`
+  query dynamicProducts($delayDeferred: Int!, $someArgument: String) {
+    products(someArgument: $someArgument) {
+      id
+      title
+      ...RatingFragment @defer
+    }
+    env
+  }
+
+  ${RATING_FRAGMENT}
+`;
diff --git a/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts b/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts
index 88b95b08..6feb8d12 100644
--- a/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts
+++ b/packages/client-react-streaming/src/DataTransportAbstraction/hooks.ts
@@ -33,6 +33,9 @@ export const hookWrappers: HookWrappers = {
       return orig_useQueryRefHandlers(useWrapTransportedQueryRef(queryRef));
     }, []);
   },
+  useSuspenseFragment(orig_useSuspenseFragment) {
+    return wrap(orig_useSuspenseFragment, ["data"]);
+  },
 };
 
 function wrap<T extends (...args: any[]) => any>(