From 8607caed8250de876599654fd461cfe5563a5e9a Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Wed, 7 Aug 2024 11:07:11 -0600
Subject: [PATCH] Allow partial return data from modifier functions

---
 src/cache/core/__tests__/cache.ts | 31 +++++++++++++++++++++++++++++++
 src/cache/core/types/common.ts    |  3 ++-
 src/cache/inmemory/entityStore.ts |  2 +-
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/cache/core/__tests__/cache.ts b/src/cache/core/__tests__/cache.ts
index 45d920be35d..accee868452 100644
--- a/src/cache/core/__tests__/cache.ts
+++ b/src/cache/core/__tests__/cache.ts
@@ -540,5 +540,36 @@ describe.skip("Cache type tests", () => {
         },
       });
     });
+
+    test("Allows partial return data", () => {
+      const cache = new TestCache();
+      cache.modify<{
+        union: Array<
+          | { __typename: "Type1"; a: string; c: { foo: string } }
+          | { __typename: "Type2"; b: string; d: { bar: number } }
+        >;
+      }>({
+        fields: {
+          union(field) {
+            expectTypeOf(field).toEqualTypeOf<
+              ReadonlyArray<
+                | Reference
+                | {
+                    __typename: "Type1";
+                    a: string;
+                    c: { foo: string };
+                  }
+                | {
+                    __typename: "Type2";
+                    b: string;
+                    d: { bar: number };
+                  }
+              >
+            >();
+            return [{ __typename: "Type1", a: "foo" }];
+          },
+        },
+      });
+    });
   });
 });
diff --git a/src/cache/core/types/common.ts b/src/cache/core/types/common.ts
index d3a200ac636..c39af692b6d 100644
--- a/src/cache/core/types/common.ts
+++ b/src/cache/core/types/common.ts
@@ -6,6 +6,7 @@ import type {
   StoreValue,
   isReference,
   AsStoreObject,
+  DeepPartial,
 } from "../../../utilities/index.js";
 
 import type { StorageType } from "../../inmemory/policies.js";
@@ -107,7 +108,7 @@ export type ModifierDetails = {
 export type Modifier<T> = (
   value: T,
   details: ModifierDetails
-) => T | DeleteModifier | InvalidateModifier | undefined;
+) => DeepPartial<T> | DeleteModifier | InvalidateModifier | undefined;
 
 type StoreObjectValueMaybeReference<StoreVal> =
   StoreVal extends Array<Record<string, any>> ?
diff --git a/src/cache/inmemory/entityStore.ts b/src/cache/inmemory/entityStore.ts
index c0582843fdc..6c85d1f5866 100644
--- a/src/cache/inmemory/entityStore.ts
+++ b/src/cache/inmemory/entityStore.ts
@@ -257,7 +257,7 @@ export abstract class EntityStore implements NormalizedCache {
             if (newValue !== fieldValue) {
               changedFields[storeFieldName] = newValue;
               needToMerge = true;
-              fieldValue = newValue;
+              fieldValue = newValue as StoreValue;
 
               if (__DEV__) {
                 const checkReference = (ref: Reference) => {