From f122e0b581187e732b28b5634ba65cb1e15912ee Mon Sep 17 00:00:00 2001
From: Alessia Bellisario <github@bellisar.io>
Date: Fri, 5 Jan 2024 12:55:11 -0500
Subject: [PATCH] feat: extract watchFragment into separate method on
 ApolloCache

---
 src/cache/core/cache.ts        | 75 +++++++++++++++++++++++++++++++++-
 src/react/hooks/useFragment.ts | 14 +++----
 2 files changed, 79 insertions(+), 10 deletions(-)

diff --git a/src/cache/core/cache.ts b/src/cache/core/cache.ts
index a0fd1778bdc..1fa039d8bf3 100644
--- a/src/cache/core/cache.ts
+++ b/src/cache/core/cache.ts
@@ -1,19 +1,52 @@
 import type { DocumentNode } from "graphql";
 import { wrap } from "optimism";
 
-import type { StoreObject, Reference } from "../../utilities/index.js";
+import type {
+  StoreObject,
+  Reference,
+  DeepPartial,
+} from "../../utilities/index.js";
 import {
+  Observable,
   cacheSizes,
   defaultCacheSizes,
   getFragmentQueryDocument,
+  mergeDeepArray,
 } from "../../utilities/index.js";
 import type { DataProxy } from "./types/DataProxy.js";
 import type { Cache } from "./types/Cache.js";
 import { WeakCache } from "@wry/caches";
 import { getApolloCacheMemoryInternals } from "../../utilities/caching/getMemoryInternals.js";
+import type {
+  OperationVariables,
+  TypedDocumentNode,
+} from "../../core/types.js";
+import { equal } from "@wry/equality";
+import type { MissingTree } from "./types/common.js";
 
 export type Transaction<T> = (c: ApolloCache<T>) => void;
 
+// todo: confirm we can start with this limited set of options
+export interface WatchFragmentOptions<TData, TVars> {
+  fragment: DocumentNode | TypedDocumentNode<TData, TVars>;
+  from: StoreObject | Reference | string;
+  fragmentName?: string;
+  optimistic?: boolean;
+}
+
+// todo: refactor duplicated type, see UseFragmentResult
+export type WatchFragmentResult<TData> =
+  | {
+      data: TData;
+      complete: true;
+      missing?: undefined;
+    }
+  | {
+      data: DeepPartial<TData>;
+      complete: false;
+      missing?: MissingTree | undefined;
+    };
+
 export abstract class ApolloCache<TSerialized> implements DataProxy {
   public readonly assumeImmutableResults: boolean = false;
 
@@ -141,6 +174,46 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
     });
   }
 
+  public watchFragment<TData = any, TVars = OperationVariables>(
+    options: WatchFragmentOptions<TData, TVars>
+  ): Observable<WatchFragmentResult<TData>> {
+    const { fragment, fragmentName, from, optimistic = true } = options;
+
+    const diffOptions: Cache.DiffOptions<TData, TVars> = {
+      returnPartialData: true,
+      id: typeof from === "string" ? from : this.identify(from),
+      query: this.getFragmentDoc(fragment, fragmentName),
+      optimistic,
+    };
+
+    let latestDiff = this.diff<TData>(diffOptions);
+
+    return new Observable((observer) => {
+      return this.watch<TData, TVars>({
+        ...diffOptions,
+        immediate: true, // is this necessary? tests pass without it
+        query: this.getFragmentDoc(fragment, fragmentName),
+        callback(diff) {
+          const result = {
+            data: diff.result as DeepPartial<TData>,
+            complete: !!diff.complete,
+          } as WatchFragmentResult<TData>;
+
+          if (diff.missing) {
+            result.missing = mergeDeepArray(
+              diff.missing.map((error) => error.missing)
+            );
+          }
+
+          if (!equal(diff, latestDiff)) {
+            latestDiff = diff;
+            observer.next(result);
+          }
+        },
+      });
+    });
+  }
+
   // Make sure we compute the same (===) fragment query document every
   // time we receive the same fragment in readFragment.
   private getFragmentDoc = wrap(getFragmentQueryDocument, {
diff --git a/src/react/hooks/useFragment.ts b/src/react/hooks/useFragment.ts
index 8764fc0cf39..5265413e037 100644
--- a/src/react/hooks/useFragment.ts
+++ b/src/react/hooks/useFragment.ts
@@ -73,18 +73,14 @@ export function useFragment<TData = any, TVars = OperationVariables>(
   return useSyncExternalStore(
     (forceUpdate) => {
       let lastTimeout = 0;
-      const unsubcribe = cache.watch({
-        ...diffOptions,
-        immediate: true,
-        callback(diff) {
-          if (!equal(diff, latestDiff)) {
-            resultRef.current = diffToResult((latestDiff = diff));
-            lastTimeout = setTimeout(forceUpdate) as any;
-          }
+      const watchedFragment = cache.watchFragment(options).subscribe({
+        next: (result) => {
+          resultRef.current = result;
+          lastTimeout = setTimeout(forceUpdate) as any;
         },
       });
       return () => {
-        unsubcribe();
+        watchedFragment.unsubscribe();
         clearTimeout(lastTimeout);
       };
     },