From 62c6a1dbbc55298fe454f1d2b88d69948c6f6dd9 Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Tue, 27 Aug 2024 10:05:54 -0600
Subject: [PATCH] Move core fragment masking logic back to cache

---
 src/cache/core/cache.ts  | 19 +++++++++++--------
 src/core/ApolloClient.ts | 39 ++++-----------------------------------
 src/core/QueryManager.ts | 16 +---------------
 3 files changed, 16 insertions(+), 58 deletions(-)

diff --git a/src/cache/core/cache.ts b/src/cache/core/cache.ts
index d8ce265787a..7cc13df2f94 100644
--- a/src/cache/core/cache.ts
+++ b/src/cache/core/cache.ts
@@ -29,6 +29,7 @@ import type {
 import type { MissingTree } from "./types/common.js";
 import { equalByQuery } from "../../core/equalByQuery.js";
 import { invariant } from "../../utilities/globals/index.js";
+import { maskFragment } from "../../core/masking.js";
 
 export type Transaction<T> = (c: ApolloCache<T>) => void;
 
@@ -252,6 +253,7 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
     } = options;
     const query = this.getFragmentDoc(fragment, fragmentName);
     const id = typeof from === "string" ? from : this.identify(from);
+    const dataMasking = !!(options as any)[Symbol.for("apollo.dataMasking")];
 
     if (__DEV__) {
       const actualFragmentName =
@@ -279,21 +281,22 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
       return this.watch<TData, TVars>({
         ...diffOptions,
         immediate: true,
-        callback(diff) {
+        callback: (diff) => {
+          const data =
+            dataMasking ?
+              maskFragment(diff.result, fragment, this, fragmentName)
+            : diff.result;
+
           if (
             // Always ensure we deliver the first result
             latestDiff &&
-            equalByQuery(
-              query,
-              { data: latestDiff?.result },
-              { data: diff.result }
-            )
+            equalByQuery(query, { data: latestDiff?.result }, { data })
           ) {
             return;
           }
 
           const result = {
-            data: diff.result as DeepPartial<TData>,
+            data,
             complete: !!diff.complete,
           } as WatchFragmentResult<TData>;
 
@@ -303,7 +306,7 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
             );
           }
 
-          latestDiff = diff;
+          latestDiff = { ...diff, result: data };
           observer.next(result);
         },
       });
diff --git a/src/core/ApolloClient.ts b/src/core/ApolloClient.ts
index fd6127e930b..765ebd512c6 100644
--- a/src/core/ApolloClient.ts
+++ b/src/core/ApolloClient.ts
@@ -6,7 +6,7 @@ import type { FetchResult, GraphQLRequest } from "../link/core/index.js";
 import { ApolloLink, execute } from "../link/core/index.js";
 import type { ApolloCache, DataProxy, Reference } from "../cache/index.js";
 import type { DocumentTransform } from "../utilities/index.js";
-import { Observable } from "../utilities/index.js";
+import type { Observable } from "../utilities/index.js";
 import { version } from "../version.js";
 import type { UriFunction } from "../link/http/index.js";
 import { HttpLink } from "../link/http/index.js";
@@ -164,7 +164,6 @@ import type {
   WatchFragmentOptions,
   WatchFragmentResult,
 } from "../cache/core/cache.js";
-import { equalByQuery } from "./equalByQuery.js";
 export { mergeOptions };
 
 /**
@@ -548,39 +547,9 @@ export class ApolloClient<TCacheShape> implements DataProxy {
   >(
     options: WatchFragmentOptions<TFragmentData, TVariables>
   ): Observable<WatchFragmentResult<TFragmentData>> {
-    const { fragment, fragmentName } = options;
-
-    const observable = this.cache.watchFragment(options);
-    let latestResult: WatchFragmentResult<TFragmentData> | undefined;
-
-    return new Observable((observer) => {
-      const subscription = observable.subscribe({
-        next: (result) => {
-          result.data = this.queryManager.maskFragment({
-            fragment,
-            fragmentName,
-            data: result.data,
-          });
-
-          if (
-            latestResult &&
-            equalByQuery(
-              this.cache["getFragmentDoc"](fragment, fragmentName),
-              { data: latestResult.data },
-              { data: result.data }
-            )
-          ) {
-            return;
-          }
-
-          latestResult = result;
-          observer.next(result);
-        },
-        complete: observer.complete.bind(observer),
-        error: observer.error.bind(observer),
-      });
-
-      return () => subscription.unsubscribe();
+    return this.cache.watchFragment({
+      ...options,
+      [Symbol.for("apollo.dataMasking")]: this.queryManager.dataMasking,
     });
   }
 
diff --git a/src/core/QueryManager.ts b/src/core/QueryManager.ts
index 07e32e94abe..523acdcba01 100644
--- a/src/core/QueryManager.ts
+++ b/src/core/QueryManager.ts
@@ -104,13 +104,7 @@ interface TransformCacheEntry {
 import type { DefaultOptions } from "./ApolloClient.js";
 import { Trie } from "@wry/trie";
 import { AutoCleanedWeakCache, cacheSizes } from "../utilities/index.js";
-import { maskFragment, maskOperation } from "./masking.js";
-
-interface MaskFragmentOptions<TData> {
-  fragment: DocumentNode;
-  data: TData;
-  fragmentName?: string;
-}
+import { maskOperation } from "./masking.js";
 
 interface MaskOperationOptions<TData> {
   document: DocumentNode;
@@ -1538,14 +1532,6 @@ export class QueryManager<TStore> {
     return this.dataMasking ? maskOperation(data, document, this.cache) : data;
   }
 
-  public maskFragment<TData = unknown>(options: MaskFragmentOptions<TData>) {
-    const { data, fragment, fragmentName } = options;
-
-    return this.dataMasking ?
-        maskFragment(data, fragment, this.cache, fragmentName)
-      : data;
-  }
-
   private fetchQueryByPolicy<TData, TVars extends OperationVariables>(
     queryInfo: QueryInfo,
     {