From 5d0f2adba2e1a177cef34d2428c44c073f32bef8 Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Mon, 1 Jul 2024 18:40:59 -0600
Subject: [PATCH] Allow disabling warnings for masked field access with @unmask

---
 src/__tests__/client.ts | 64 +++++++++++++++++++++++++++++++++++++++++
 src/core/masking.ts     | 35 ++++++++++++++++------
 2 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/src/__tests__/client.ts b/src/__tests__/client.ts
index 766da6bb53a..c61ca16580b 100644
--- a/src/__tests__/client.ts
+++ b/src/__tests__/client.ts
@@ -7478,6 +7478,70 @@ describe("data masking", () => {
       expect(consoleSpy.warn).toHaveBeenCalledTimes(1);
     }
   });
+
+  it("allows disabling warnings when accessing a fragmented field while using @unmask", async () => {
+    using consoleSpy = spyOnConsole("warn");
+
+    interface Query {
+      currentUser: {
+        __typename: "User";
+        id: number;
+        name: string;
+        age: number;
+      };
+    }
+
+    const query: TypedDocumentNode<Query, never> = gql`
+      query UnmaskedQuery @unmask(warnOnFieldAccess: false) {
+        currentUser {
+          id
+          name
+          ...UserFields
+        }
+      }
+
+      fragment UserFields on User {
+        age
+        name
+      }
+    `;
+
+    const mocks = [
+      {
+        request: { query },
+        result: {
+          data: {
+            currentUser: {
+              __typename: "User",
+              id: 1,
+              name: "Test User",
+              age: 34,
+            },
+          },
+        },
+        delay: 20,
+      },
+    ];
+
+    const client = new ApolloClient({
+      dataMasking: true,
+      cache: new InMemoryCache(),
+      link: new MockLink(mocks),
+    });
+
+    const observable = client.watchQuery({ query });
+    const stream = new ObservableStream(observable);
+
+    {
+      const { data } = await stream.takeNext();
+      data.currentUser.__typename;
+      data.currentUser.id;
+      data.currentUser.name;
+      data.currentUser.age;
+
+      expect(consoleSpy.warn).not.toHaveBeenCalled();
+    }
+  });
 });
 
 function clientRoundtrip(
diff --git a/src/core/masking.ts b/src/core/masking.ts
index d7511b625af..803c7297f9a 100644
--- a/src/core/masking.ts
+++ b/src/core/masking.ts
@@ -27,12 +27,14 @@ export function maskQuery<TData = unknown>(
 ): TData {
   const definition = getMainDefinition(document);
   const fragmentMap = createFragmentMap(getFragmentDefinitions(document));
+  const [isUnmasked, { warnOnFieldAccess }] = isUnmaskedDocument(document);
   const [masked, changed] = maskSelectionSet(
     data,
     definition.selectionSet,
     matchesFragment,
     fragmentMap,
-    isUnmaskedDocument(document)
+    isUnmasked,
+    warnOnFieldAccess
   );
 
   return changed ? masked : data;
@@ -75,7 +77,8 @@ export function maskFragment<TData = unknown>(
     fragment.selectionSet,
     matchesFragment,
     fragmentMap,
-    false
+    false,
+    true
   );
 
   return changed ? masked : data;
@@ -86,7 +89,8 @@ function maskSelectionSet(
   selectionSet: SelectionSetNode,
   matchesFragment: MatchesFragmentFn,
   fragmentMap: FragmentMap,
-  isUnmasked: boolean
+  isUnmasked: boolean,
+  warnOnFieldAccess: boolean
 ): [data: any, changed: boolean] {
   if (Array.isArray(data)) {
     let changed = false;
@@ -97,7 +101,8 @@ function maskSelectionSet(
         selectionSet,
         matchesFragment,
         fragmentMap,
-        isUnmasked
+        isUnmasked,
+        warnOnFieldAccess
       );
       changed ||= itemChanged;
 
@@ -132,7 +137,8 @@ function maskSelectionSet(
               childSelectionSet,
               matchesFragment,
               fragmentMap,
-              isUnmasked
+              isUnmasked,
+              warnOnFieldAccess
             );
 
             if (childChanged) {
@@ -156,7 +162,8 @@ function maskSelectionSet(
             selection.selectionSet,
             matchesFragment,
             fragmentMap,
-            isUnmasked
+            isUnmasked,
+            warnOnFieldAccess
           );
 
           return [
@@ -172,7 +179,12 @@ function maskSelectionSet(
 
           return [
             isUnmasked ?
-              addAccessorWarnings(memo, data, fragment.selectionSet)
+              addAccessorWarnings(
+                memo,
+                data,
+                fragment.selectionSet,
+                warnOnFieldAccess
+              )
             : memo,
             true,
           ];
@@ -185,7 +197,8 @@ function maskSelectionSet(
 function addAccessorWarnings(
   memo: Record<string, unknown>,
   parent: Record<string, unknown>,
-  selectionSetNode: SelectionSetNode
+  selectionSetNode: SelectionSetNode,
+  warnOnFieldAccess: boolean
 ) {
   selectionSetNode.selections.forEach((selection) => {
     switch (selection.kind) {
@@ -196,7 +209,11 @@ function addAccessorWarnings(
           return;
         }
 
-        return addAccessorWarning(memo, parent[keyName], keyName);
+        if (warnOnFieldAccess) {
+          return addAccessorWarning(memo, parent[keyName], keyName);
+        } else {
+          memo[keyName] = parent[keyName];
+        }
       }
     }
   });