Preserve the obserable from cache and pass a transform function

Author: jerelmiller

src/cache/core/cache.ts

@@ -439,6 +439,10 @@ export abstract class ApolloCache {
       fragment,
       fragmentName
     ) as TypedDocumentNode<TData, TVariables>;
+    const transform: (
+      result: ApolloCache.WatchFragmentResult<TData>
+    ) => ApolloCache.WatchFragmentResult<TData> =
+      (options as any)[Symbol.for("apollo.transform")] ?? ((v) => v);
 
     const fromArray = Array.isArray(from) ? from : [from];
 
@@ -486,15 +490,17 @@ export abstract class ApolloCache {
         // without the need for a similar fallback since watchFragment with
         // arrays is new functionality in v4.1.
         transform: (result) =>
-          from === null ? result : { ...result, data: result.data ?? {} },
+          transform(
+            from === null ? result : { ...result, data: result.data ?? {} }
+          ),
       });
     }
 
-    let currentResult: ApolloCache.WatchFragmentResult<any>;
+    let currentResult: ApolloCache.WatchFragmentResult<TData>;
     function toResult(
       results: Array<ApolloCache.WatchFragmentResult<Unmasked<TData> | null>>
     ): ApolloCache.WatchFragmentResult<any> {
-      const result = results.reduce(
+      let result = results.reduce(
         (finalResult, res, idx) => {
           const result = res as ApolloCache.WatchFragmentResult<TData>;
 
@@ -513,9 +519,11 @@ export abstract class ApolloCache {
           data: [],
           dataState: "complete",
           complete: true,
-        } as ApolloCache.WatchFragmentResult<any>
+        } as ApolloCache.WatchFragmentResult<TData>
       );
 
+      result = transform(result);
+
       if (!equal(currentResult, result)) {
         currentResult = result;
       }
@@ -617,7 +625,15 @@ export abstract class ApolloCache {
             this.onAfterBroadcast(() => {
               const result = transform(toWatchFragmentResult(diff));
 
-              if (!equal(currentResult, result)) {
+              if (
+                !currentResult ||
+                !equalByQuery(
+                  fragmentQuery,
+                  { data: currentResult.data },
+                  { data: result.data },
+                  options.variables
+                )
+              ) {
                 currentResult = result;
               }
 
@@ -632,14 +648,7 @@ export abstract class ApolloCache {
           this.fragmentWatches.removeArray(cacheKey);
         };
       }).pipe(
-        distinctUntilChanged((previous, current) =>
-          equalByQuery(
-            fragmentQuery,
-            { data: previous.data },
-            { data: current.data },
-            options.variables
-          )
-        ),
+        distinctUntilChanged(),
         share({
           connector: () => new ReplaySubject(1),
           // debounce so a synchronous unsubscribe+resubscribe doesn't tear down the watch and create a new one

src/core/ApolloClient.ts

@@ -1187,10 +1187,6 @@ export class ApolloClient {
     | ApolloClient.ObservableFragment<TData>
     | ApolloClient.ObservableFragment<Array<TData>> {
     const dataMasking = this.queryManager.dataMasking;
-    const observable = this.cache.watchFragment({
-      ...options,
-      fragment: this.transform(options.fragment, dataMasking),
-    });
 
     const mask = (
       result: ApolloClient.WatchFragmentResult<any>
@@ -1215,21 +1211,13 @@ export class ApolloClient {
       return result;
     };
 
-    let currentResult: ApolloClient.WatchFragmentResult<any>;
-    let stableMaskedResult: ApolloClient.WatchFragmentResult<any>;
-
-    return Object.assign(observable.pipe(map(mask)), {
-      getCurrentResult: () => {
-        const result = observable.getCurrentResult();
-
-        if (result !== currentResult) {
-          currentResult = result as any;
-          stableMaskedResult = mask(currentResult);
-        }
+    const observable = this.cache.watchFragment({
+      ...options,
+      fragment: this.transform(options.fragment, dataMasking),
+      [Symbol.for("apollo.transform")]: mask,
+    });
 
-        return stableMaskedResult;
-      },
-    }) as ApolloClient.ObservableFragment<any>;
+    return observable as ApolloClient.ObservableFragment<any>;
   }
 
   /**