From 67c1e1cb18b242644a36b832607b2abaea4a6714 Mon Sep 17 00:00:00 2001
From: Alessia Bellisario <github@bellisar.io>
Date: Fri, 12 Jan 2024 16:46:07 -0500
Subject: [PATCH] adjustments based on review comments

---
 src/__tests__/ApolloClient.ts  |  2 +-
 src/cache/core/cache.ts        | 22 +++++++++++++++++++---
 src/core/ApolloClient.ts       | 13 +++++++------
 src/react/hooks/useFragment.ts | 20 ++------------------
 4 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/src/__tests__/ApolloClient.ts b/src/__tests__/ApolloClient.ts
index c6b7a3cae35..35511eda9b9 100644
--- a/src/__tests__/ApolloClient.ts
+++ b/src/__tests__/ApolloClient.ts
@@ -2420,7 +2420,7 @@ describe("ApolloClient", () => {
 
         const handleNext = jest.fn();
 
-        observable.subscribe(handleNext);
+        queryObs.subscribe(handleNext);
 
         cache.writeQuery({
           query,
diff --git a/src/cache/core/cache.ts b/src/cache/core/cache.ts
index 566add47d61..68b4a67f92c 100644
--- a/src/cache/core/cache.ts
+++ b/src/cache/core/cache.ts
@@ -30,8 +30,23 @@ export type Transaction<T> = (c: ApolloCache<T>) => void;
 export interface WatchFragmentOptions<TData, TVars> {
   fragment: DocumentNode | TypedDocumentNode<TData, TVars>;
   from: StoreObject | Reference | string;
+  /**
+   * Any variables that the GraphQL query may depend on.
+   */
+  variables?: TVars;
   fragmentName?: string;
   optimistic?: boolean;
+  /**
+   * @deprecated
+   * Using `canonizeResults` can result in memory leaks so we generally do not
+   * recommend using this option anymore.
+   * A future version of Apollo Client will contain a similar feature.
+   *
+   * Whether to canonize cache results before returning them. Canonization
+   * takes some extra time, but it speeds up future deep equality comparisons.
+   * Defaults to false.
+   */
+  canonizeResults?: boolean;
 }
 
 // todo: refactor duplicated type, see UseFragmentResult
@@ -39,12 +54,12 @@ export type WatchFragmentResult<TData> =
   | {
       data: TData;
       complete: true;
-      missing?: undefined;
+      missing?: never;
     }
   | {
       data: DeepPartial<TData>;
       complete: false;
-      missing?: MissingTree | undefined;
+      missing: MissingTree;
     };
 
 export abstract class ApolloCache<TSerialized> implements DataProxy {
@@ -186,7 +201,8 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
       optimistic,
     };
 
-    let latestDiff: DataProxy.DiffResult<TData> | undefined = undefined;
+    // let latestDiff: DataProxy.DiffResult<TData> | undefined = undefined;
+    let latestDiff: DataProxy.DiffResult<TData> = this.diff<TData>(diffOptions);
 
     return new Observable((observer) => {
       return this.watch<TData, TVars>({
diff --git a/src/core/ApolloClient.ts b/src/core/ApolloClient.ts
index 1e1ef836cc6..84ee3626e95 100644
--- a/src/core/ApolloClient.ts
+++ b/src/core/ApolloClient.ts
@@ -471,8 +471,8 @@ export class ApolloClient<TCacheShape> implements DataProxy {
 
   /**
    * Watches the cache store of the fragment according to the options specified
-   * and returns an {@link ObservableQuery}. We can subscribe to this
-   * {@link ObservableQuery} and receive updated results through a GraphQL
+   * and returns an {@link Observable}. We can subscribe to this
+   * {@link Observable} and receive updated results through a GraphQL
    * observer when the cache store changes.
    *
    * You must pass in a GraphQL document with a single fragment or a document
@@ -485,10 +485,11 @@ export class ApolloClient<TCacheShape> implements DataProxy {
    * to optimistic updates.
    */
 
-  public watchFragment<T = any, TVariables = OperationVariables>(
-    options: WatchFragmentOptions<T, TVariables>
-  ) {
-    return this.cache.watchFragment<T, TVariables>(options);
+  public watchFragment<
+    TFragmentData = unknown,
+    TVariables = OperationVariables,
+  >(options: WatchFragmentOptions<TFragmentData, TVariables>) {
+    return this.cache.watchFragment<TFragmentData, TVariables>(options);
   }
 
   /**
diff --git a/src/react/hooks/useFragment.ts b/src/react/hooks/useFragment.ts
index d487df2041f..a42987ab295 100644
--- a/src/react/hooks/useFragment.ts
+++ b/src/react/hooks/useFragment.ts
@@ -1,5 +1,4 @@
 import * as React from "rehackt";
-import { equal } from "@wry/equality";
 
 import type { DeepPartial } from "../../utilities/index.js";
 import { mergeDeepArray } from "../../utilities/index.js";
@@ -76,31 +75,16 @@ export function useFragment<TData = any, TVars = OperationVariables>(
     React.useCallback(
       (forceUpdate) => {
         let lastTimeout = 0;
-        const watchedFragment = cache.watchFragment(options).subscribe({
+        const subscription = cache.watchFragment(options).subscribe({
           next: (result) => {
             resultRef.current = result;
             lastTimeout = setTimeout(forceUpdate) as any;
           },
         });
         return () => {
-          watchedFragment.unsubscribe();
+          subscription.unsubscribe();
           clearTimeout(lastTimeout);
         };
-        // const unsubscribe = cache.watch({
-        //   ...diffOptions,
-        //   immediate: true,
-        //   callback(diff) {
-        //     if (!equal(diff.result, resultRef.current.data)) {
-        //       resultRef.current = diffToResult(diff);
-        //       // If we get another update before we've re-rendered, bail out of
-        //       // the update and try again. This ensures that the relative timing
-        //       // between useQuery and useFragment stays roughly the same as
-        //       // fixed in https://github.com/apollographql/apollo-client/pull/11083
-        //       clearTimeout(lastTimeout);
-        //       lastTimeout = setTimeout(forceUpdate) as any;
-        //     }
-        //   },
-        // });
       },
       [cache, diffOptions]
     ),