Type safe attempt that is backward compatible

src/__tests__/ApolloClient.ts

@@ -3219,16 +3219,24 @@ describe("ApolloClient", () => {
       >();
 
       observableQuery.updateQuery((previousData) => {
-        expectTypeOf(previousData).toMatchTypeOf<UnmaskedQuery | undefined>();
-        expectTypeOf(previousData).not.toMatchTypeOf<Query>();
+        expectTypeOf(previousData).toMatchTypeOf<Partial<UnmaskedQuery>>();
+        if (previousData.complete) {
+          expectTypeOf(previousData).toMatchTypeOf<UnmaskedQuery>();
+          expectTypeOf(previousData).not.toMatchTypeOf<Query>();
+        } else {
+          expectTypeOf(previousData).toMatchTypeOf<Partial<UnmaskedQuery>>();
+          expectTypeOf(previousData).not.toMatchTypeOf<Query>();
+        }
 
-        return {} as UnmaskedQuery;
+        return undefined;
       });
 
       observableQuery.subscribeToMore({
         document: subscription,
         updateQuery(queryData, { subscriptionData }) {
-          expectTypeOf(queryData).toMatchTypeOf<UnmaskedQuery | undefined>();
+          expectTypeOf(queryData).toMatchTypeOf<
+            UnmaskedQuery | Partial<UnmaskedQuery>
+          >();
           expectTypeOf(queryData).not.toMatchTypeOf<Query>();
 
           expectTypeOf(

src/__tests__/dataMasking.ts

@@ -1892,15 +1892,12 @@ describe("client.watchQuery", () => {
 
     const updateQuery: Parameters<typeof observable.updateQuery>[0] = jest.fn(
       (previousResult) => {
-        return {
-          user: {
-            __typename: "User",
-            id: 1,
-            age: 30,
-            ...previousResult?.user,
-            name: "User (updated)",
-          },
-        };
+        expect(previousResult.complete).toBe(true);
+        // Type guard
+        if (!previousResult.complete) {
+          return undefined;
+        }
+        return { user: { ...previousResult.user, name: "User (updated)" } };
       }
     );
 

src/__tests__/subscribeToMore.ts

@@ -241,10 +241,14 @@ describe("subscribeToMore", () => {
         }
       `,
       updateQuery: (prev, { subscriptionData }) => {
-        expect(prev).toBeTruthy();
-        expect(prev!.entry).not.toContainEqual(nextMutation);
+        expect(prev.complete).toBe(true);
+        // Type guard
+        if (!prev.complete) {
+          return undefined;
+        }
+        expect(prev.entry).not.toContainEqual(nextMutation);
         return {
-          entry: [...prev!.entry, { value: subscriptionData.data.name }],
+          entry: [...prev.entry, { value: subscriptionData.data.name }],
         };
       },
     });

src/core/ObservableQuery.ts

@@ -31,6 +31,7 @@ import type {
   SubscribeToMoreOptions,
   NextFetchPolicyContext,
   WatchQueryFetchPolicy,
+  UpdateQueryMapFn,
 } from "./watchQueryOptions.js";
 import type { QueryInfo } from "./QueryInfo.js";
 import type { MissingFieldError } from "../cache/index.js";
@@ -619,7 +620,8 @@ Did you mean to call refetch(variables) instead of refetch({ variables })?`,
     options: SubscribeToMoreOptions<
       TData,
       TSubscriptionVariables,
-      TSubscriptionData
+      TSubscriptionData,
+      TVariables
     >
   ) {
     const subscription = this.queryManager
@@ -632,12 +634,12 @@ Did you mean to call refetch(variables) instead of refetch({ variables })?`,
         next: (subscriptionData: { data: Unmasked<TSubscriptionData> }) => {
           const { updateQuery } = options;
           if (updateQuery) {
-            this.updateQuery<TSubscriptionVariables>(
-              (previous, { variables }) =>
-                updateQuery(previous, {
-                  subscriptionData,
-                  variables,
-                })
+            this.updateQuery((previous, { variables }) =>
+              updateQuery(previous, {
+                subscriptionData,
+                subscriptionVariables: options.variables,
+                variables,
+              })
             );
           }
         },
@@ -722,28 +724,35 @@ Did you mean to call refetch(variables) instead of refetch({ variables })?`,
    *
    * See [using updateQuery and updateFragment](https://www.apollographql.com/docs/react/caching/cache-interaction/#using-updatequery-and-updatefragment) for additional information.
    */
-  public updateQuery<TVars extends OperationVariables = TVariables>(
-    mapFn: (
-      previousQueryResult: Unmasked<TData> | undefined,
-      options: Pick<WatchQueryOptions<TVars, TData>, "variables">
-    ) => Unmasked<TData> | undefined
-  ): void {
+  public updateQuery(mapFn: UpdateQueryMapFn<TData, TVariables>): void {
     const { queryManager } = this;
-    const { result } = queryManager.cache.diff<Unmasked<TData>>({
+    const { result, complete } = queryManager.cache.diff<Unmasked<TData>>({
       query: this.options.query,
       variables: this.variables,
       returnPartialData: true,
       optimistic: false,
     });
 
-    const newResult = mapFn(result, {
+    const completeSymbol = Symbol("complete");
+    if (complete) {
+      Object.defineProperty(result, "complete", {
+        value: completeSymbol,
+        writable: false,
+        enumerable: false,
+        configurable: false,
+      });
+    }
+
+    const newResult = mapFn(result as any, {
       variables: (this as any).variables,
     });
 
     if (newResult) {
       queryManager.cache.writeQuery({
         query: this.options.query,
-        data: newResult,
+        // Spread the new result to ensure that the complete property is not
+        // included in the result since it's not enumerable.
+        data: newResult === result ? { ...newResult } : newResult,
         variables: this.variables,
       });
 

src/core/watchQueryOptions.ts

@@ -164,31 +164,79 @@ export interface FetchMoreQueryOptions<TVariables, TData = any> {
   context?: DefaultContext;
 }
 
-export type UpdateQueryFn<
+export interface UpdateQueryFn<
+  TData,
+  TVars extends OperationVariables,
+  TOptions = {},
+> {
+  (mapFn: UpdateQueryMapFn<TData, TVars, TOptions>): void;
+}
+
+export interface UpdateQueryMapFn<
+  TData = any,
+  TVars extends OperationVariables = OperationVariables,
+  TOptions = {},
+> {
+  (
+    previousQueryResult: // DeepPartial is more accurate, but causes typescript to stop on
+    // "Type instantiation is excessively deep and possibly infinite."
+    // | (DeepPartial<Unmasked<TData>> & { complete?: undefined })
+    | (Partial<Unmasked<TData>> & { complete?: undefined })
+      | (Unmasked<TData> & { complete: Symbol }),
+    options: TOptions & { variables?: TVars }
+  ): Unmasked<TData> | undefined;
+}
+
+export type SubscribeToMoreUpdateQueryFn<
   TData = any,
-  TSubscriptionVariables = OperationVariables,
+  TVars extends OperationVariables = OperationVariables,
+  TSubscriptionVariables extends OperationVariables = TVars,
   TSubscriptionData = TData,
-> = (
-  previousQueryResult: Unmasked<TData> | undefined,
-  options: {
+> = UpdateQueryMapFn<
+  TData,
+  TVars,
+  {
     subscriptionData: { data: Unmasked<TSubscriptionData> };
-    variables?: TSubscriptionVariables;
+    subscriptionVariables: TSubscriptionVariables | undefined;
   }
-) => Unmasked<TData> | undefined;
+>;
 
-export type SubscribeToMoreOptions<
+export interface SubscribeToMoreOptions<
   TData = any,
-  TSubscriptionVariables = OperationVariables,
+  TSubscriptionVariables extends OperationVariables = OperationVariables,
   TSubscriptionData = TData,
-> = {
+  TVars extends OperationVariables = TSubscriptionVariables,
+> {
   document:
     | DocumentNode
     | TypedDocumentNode<TSubscriptionData, TSubscriptionVariables>;
   variables?: TSubscriptionVariables;
-  updateQuery?: UpdateQueryFn<TData, TSubscriptionVariables, TSubscriptionData>;
+  updateQuery?: SubscribeToMoreUpdateQueryFn<
+    TData,
+    TVars,
+    TSubscriptionVariables,
+    TSubscriptionData
+  >;
   onError?: (error: Error) => void;
   context?: DefaultContext;
-};
+}
+
+export interface SubscribeToMoreFunction<
+  TData,
+  TVars extends OperationVariables = OperationVariables,
+> {
+  <
+    TSubscriptionData = TData,
+    TSubscriptionVariables extends OperationVariables = TVars,
+  >(
+    options: SubscribeToMoreOptions<
+      TData,
+      TSubscriptionVariables,
+      TSubscriptionData,
+      TVars
+    >
+  ): () => void;
+}
 
 export interface SubscriptionOptions<
   TVariables = OperationVariables,

src/react/hooks/__tests__/useBackgroundQuery.test.tsx

@@ -42,6 +42,7 @@ import equal from "@wry/equality";
 import {
   RefetchWritePolicy,
   SubscribeToMoreOptions,
+  SubscribeToMoreFunction,
 } from "../../../core/watchQueryOptions";
 import { skipToken } from "../constants";
 import {
@@ -57,7 +58,6 @@ import {
   spyOnConsole,
   addDelayToMocks,
 } from "../../../testing/internal";
-import { SubscribeToMoreFunction } from "../useSuspenseQuery";
 import {
   MaskedVariablesCaseData,
   setupMaskedVariablesCase,
@@ -8364,9 +8364,12 @@ describe.skip("type tests", () => {
     {
       const [, { subscribeToMore }] = useBackgroundQuery(query);
 
-      const subscription: MaskedDocumentNode<Subscription, never> = gql`
-        subscription {
-          pushLetter {
+      const subscription: MaskedDocumentNode<
+        Subscription,
+        { letterId: string }
+      > = gql`
+        subscription LetterPushed($letterId: ID!) {
+          pushLetter(letterId: $letterId) {
             id
             ...CharacterFragment
           }
@@ -8379,9 +8382,13 @@ describe.skip("type tests", () => {
 
       subscribeToMore({
         document: subscription,
-        updateQuery: (queryData, { subscriptionData }) => {
+        updateQuery: (
+          queryData,
+          { subscriptionData, subscriptionVariables, variables }
+        ) => {
           expectTypeOf(queryData).toEqualTypeOf<
-            UnmaskedVariablesCaseData | undefined
+            | (Partial<UnmaskedVariablesCaseData> & { complete?: undefined })
+            | (UnmaskedVariablesCaseData & { complete: Symbol })
           >();
           expectTypeOf(queryData).not.toEqualTypeOf<MaskedVariablesCaseData>();
 
@@ -8390,6 +8397,14 @@ describe.skip("type tests", () => {
           ).toEqualTypeOf<UnmaskedSubscription>();
           expectTypeOf(subscriptionData.data).not.toEqualTypeOf<Subscription>();
 
+          expectTypeOf(subscriptionVariables).toEqualTypeOf<
+            { letterId: string } | undefined
+          >();
+
+          expectTypeOf(variables).toEqualTypeOf<
+            VariablesCaseVariables | undefined
+          >();
+
           return {} as UnmaskedVariablesCaseData;
         },
       });
@@ -8413,9 +8428,13 @@ describe.skip("type tests", () => {
 
       subscribeToMore({
         document: subscription,
-        updateQuery: (queryData, { subscriptionData }) => {
+        updateQuery: (
+          queryData,
+          { subscriptionData, subscriptionVariables, variables }
+        ): UnmaskedVariablesCaseData => {
           expectTypeOf(queryData).toEqualTypeOf<
-            UnmaskedVariablesCaseData | undefined
+            | (Partial<UnmaskedVariablesCaseData> & { complete?: undefined })
+            | (UnmaskedVariablesCaseData & { complete: Symbol })
           >();
           expectTypeOf(queryData).not.toEqualTypeOf<MaskedVariablesCaseData>();
 
@@ -8424,7 +8443,20 @@ describe.skip("type tests", () => {
           ).toEqualTypeOf<UnmaskedSubscription>();
           expectTypeOf(subscriptionData.data).not.toEqualTypeOf<Subscription>();
 
-          return {} as UnmaskedVariablesCaseData;
+          expectTypeOf(subscriptionVariables).toEqualTypeOf<undefined>();
+
+          expectTypeOf(variables).toEqualTypeOf<
+            VariablesCaseVariables | undefined
+          >();
+
+          if (queryData.complete) {
+            expectTypeOf(queryData).toEqualTypeOf<
+              UnmaskedVariablesCaseData & { complete: Symbol }
+            >();
+            return queryData;
+          }
+          // @ts-expect-error -- The incomplete data cannot be returned
+          return queryData;
         },
       });
     }

src/react/hooks/__tests__/useLoadableQuery.test.tsx

@@ -18,6 +18,7 @@ import {
   SubscribeToMoreOptions,
   split,
 } from "../../../core";
+import { SubscribeToMoreFunction } from "../../../core/watchQueryOptions";
 import {
   MockedProvider,
   MockedProviderProps,
@@ -42,7 +43,6 @@ import { QueryRef } from "../../../react";
 import {
   FetchMoreFunction,
   RefetchFunction,
-  SubscribeToMoreFunction,
 } from "../useSuspenseQuery";
 import invariant, { InvariantError } from "ts-invariant";
 import {

src/react/hooks/__tests__/useQueryRefHandlers.test.tsx

@@ -9,6 +9,7 @@ import {
   gql,
   split,
 } from "../../../core";
+import { SubscribeToMoreFunction } from "../../../core/watchQueryOptions";
 import {
   MockLink,
   MockSubscriptionLink,
@@ -23,7 +24,6 @@ import {
 } from "../../../testing/internal";
 import { useQueryRefHandlers } from "../useQueryRefHandlers";
 import { UseReadQueryResult, useReadQuery } from "../useReadQuery";
-import type { SubscribeToMoreFunction } from "../useSuspenseQuery";
 import { Suspense } from "react";
 import { createQueryPreloader } from "../../query-preloader/createQueryPreloader";
 import userEvent from "@testing-library/user-event";

src/react/hooks/__tests__/useSuspenseQuery.test.tsx

@@ -12905,7 +12905,8 @@ describe("useSuspenseQuery", () => {
           document: subscription,
           updateQuery: (queryData, { subscriptionData }) => {
             expectTypeOf(queryData).toEqualTypeOf<
-              UnmaskedVariablesCaseData | undefined
+              | (Partial<UnmaskedVariablesCaseData> & { complete?: undefined })
+              | (UnmaskedVariablesCaseData & { complete: Symbol })
             >();
             expectTypeOf(
               queryData
@@ -12943,7 +12944,8 @@ describe("useSuspenseQuery", () => {
           document: subscription,
           updateQuery: (queryData, { subscriptionData }) => {
             expectTypeOf(queryData).toEqualTypeOf<
-              UnmaskedVariablesCaseData | undefined
+              | (Partial<UnmaskedVariablesCaseData> & { complete?: undefined })
+              | (UnmaskedVariablesCaseData & { complete: Symbol })
             >();
             expectTypeOf(
               queryData