Provide variables and context to update function

CHANGELOG.md

@@ -56,6 +56,8 @@ TBD
 - Allow `merge: true` field policy to merge `Reference` objects with non-normalized objects, and vice-versa. <br/>
   [@benjamn](https://github.com/benjamn) in [#7778](https://github.com/apollographql/apollo-client/pull/7778)
 
+- Pass `variables` and `context` to a mutation's `update` function <br/>
+  [@jcreighton](https://github.com/jcreighton) in [#7902](https://github.com/apollographql/apollo-client/pull/7902)
 ### Documentation
 TBD
 

src/core/ApolloClient.ts

@@ -12,6 +12,7 @@ import { ObservableQuery } from './ObservableQuery';
 
 import {
   ApolloQueryResult,
+  DefaultContext,
   OperationVariables,
   Resolvers,
 } from './types';
@@ -33,7 +34,7 @@ import {
 export interface DefaultOptions {
   watchQuery?: Partial<WatchQueryOptions<any, any>>;
   query?: Partial<QueryOptions<any, any>>;
-  mutate?: Partial<MutationOptions<any, any>>;
+  mutate?: Partial<MutationOptions<any, any, any>>;
 }
 
 let hasSuggestedDevtools = false;
@@ -57,13 +58,13 @@ export type ApolloClientOptions<TCacheShape> = {
   version?: string;
 };
 
-type OptionsUnion<TData, TVariables> =
+type OptionsUnion<TData, TVariables, TContext> =
   | WatchQueryOptions<TVariables, TData>
   | QueryOptions<TVariables, TData>
-  | MutationOptions<TData, TVariables>;
+  | MutationOptions<TData, TVariables, TContext>;
 
 export function mergeOptions<
-  TOptions extends OptionsUnion<any, any>
+  TOptions extends OptionsUnion<any, any, any>
 >(
   defaults: Partial<TOptions>,
   options: TOptions,
@@ -348,13 +349,18 @@ export class ApolloClient<TCacheShape> implements DataProxy {
    *
    * It takes options as an object with the following keys and values:
    */
-  public mutate<T = any, TVariables = OperationVariables>(
-    options: MutationOptions<T, TVariables>,
-  ): Promise<FetchResult<T>> {
+  public mutate<
+    TData = any,
+    TVariables = OperationVariables,
+    TContext = DefaultContext,
+    TCache extends ApolloCache<any> = ApolloCache<any>
+  >(
+    options: MutationOptions<TData, TVariables, TContext>,
+  ): Promise<FetchResult<TData>> {
     if (this.defaultOptions.mutate) {
       options = mergeOptions(this.defaultOptions.mutate, options);
     }
-    return this.queryManager.mutate<T>(options);
+    return this.queryManager.mutate<TData, TVariables, TContext, TCache>(options);
   }
 
   /**

src/core/QueryManager.ts

@@ -35,6 +35,7 @@ import { NetworkStatus, isNetworkRequestInFlight } from './networkStatus';
 import {
   ApolloQueryResult,
   OperationVariables,
+  MutationUpdaterFunction,
   ReobserveQueryCallback,
 } from './types';
 import { LocalState } from './LocalState';
@@ -55,6 +56,8 @@ interface MutationStoreValue {
   error: Error | null;
 }
 
+type UpdateQueries<TData> = MutationOptions<TData, any, any>["updateQueries"];
+
 export class QueryManager<TStore> {
   public cache: ApolloCache<TStore>;
   public link: ApolloLink;
@@ -128,7 +131,7 @@ export class QueryManager<TStore> {
     this.fetchCancelFns.clear();
   }
 
-  public async mutate<T>({
+  public async mutate<TData, TVariables, TContext, TCache extends ApolloCache<any>>({
     mutation,
     variables,
     optimisticResponse,
@@ -139,8 +142,8 @@ export class QueryManager<TStore> {
     reobserveQuery,
     errorPolicy = 'none',
     fetchPolicy,
-    context = {},
-  }: MutationOptions): Promise<FetchResult<T>> {
+    context,
+  }: MutationOptions<TData, TVariables, TContext>): Promise<FetchResult<TData>> {
     invariant(
       mutation,
       'mutation option is required. You must specify your GraphQL document in the mutation option.',
@@ -154,10 +157,10 @@ export class QueryManager<TStore> {
     const mutationId = this.generateMutationId();
     mutation = this.transform(mutation).document;
 
-    variables = this.getVariables(mutation, variables);
+    variables = this.getVariables(mutation, variables) as TVariables;
 
     if (this.transform(mutation).hasClientExports) {
-      variables = await this.localState.addExportedVariables(mutation, variables, context);
+      variables = await this.localState.addExportedVariables(mutation, variables, context) as TVariables;
     }
 
     const mutationStoreValue =
@@ -170,11 +173,17 @@ export class QueryManager<TStore> {
       } as MutationStoreValue);
 
     if (optimisticResponse) {
-      this.markMutationOptimistic<T>(optimisticResponse, {
+      this.markMutationOptimistic<
+        TData,
+        TVariables,
+        TContext,
+        TCache
+      >(optimisticResponse, {
         mutationId,
         document: mutation,
         variables,
         errorPolicy,
+        context,
         updateQueries,
         update: updateWithProxyFn,
       });
@@ -185,7 +194,7 @@ export class QueryManager<TStore> {
     const self = this;
 
     return new Promise((resolve, reject) => {
-      let storeResult: FetchResult<T> | null;
+      let storeResult: FetchResult<TData> | null;
 
       return asyncMap(
         self.getObservableFromLink(
@@ -198,7 +207,7 @@ export class QueryManager<TStore> {
           false,
         ),
 
-        (result: FetchResult<T>) => {
+        (result: FetchResult<TData>) => {
           if (graphQLResultHasError(result) && errorPolicy === 'none') {
             throw new ApolloError({
               graphQLErrors: result.errors,
@@ -218,12 +227,13 @@ export class QueryManager<TStore> {
               // mutation await any Promise that markMutationResult returns,
               // since we are returning this Promise from the asyncMap mapping
               // function.
-              return self.markMutationResult<T>({
+              return self.markMutationResult<TData, TVariables, TContext, TCache>({
                 mutationId,
                 result,
                 document: mutation,
                 variables,
                 errorPolicy,
+                context,
                 updateQueries,
                 update: updateWithProxyFn,
                 reobserveQuery,
@@ -291,18 +301,16 @@ export class QueryManager<TStore> {
     });
   }
 
-  public markMutationResult<TData>(
+  public markMutationResult<TData, TVariables, TContext, TCache extends ApolloCache<any>>(
     mutation: {
       mutationId: string;
       result: FetchResult<TData>;
       document: DocumentNode;
-      variables?: OperationVariables;
+      variables?: TVariables;
       errorPolicy: ErrorPolicy;
-      updateQueries: MutationOptions<TData>["updateQueries"],
-      update?: (
-        cache: ApolloCache<TStore>,
-        result: FetchResult<TData>,
-      ) => void;
+      context?: TContext;
+      updateQueries: UpdateQueries<TData>;
+      update?: MutationUpdaterFunction<TData, TVariables, TContext, TCache>;
       reobserveQuery?: ReobserveQueryCallback;
     },
     cache = this.cache,
@@ -364,7 +372,10 @@ export class QueryManager<TStore> {
           // a write action.
           const { update } = mutation;
           if (update) {
-            update(c, mutation.result);
+            update(c as any, mutation.result, {
+              context: mutation.context,
+              variables: mutation.variables,
+            });
           }
         },
 
@@ -389,18 +400,16 @@ export class QueryManager<TStore> {
     return Promise.resolve();
   }
 
-  public markMutationOptimistic<TData>(
+  public markMutationOptimistic<TData, TVariables, TContext, TCache extends ApolloCache<any>>(
     optimisticResponse: any,
     mutation: {
       mutationId: string;
       document: DocumentNode;
-      variables?: OperationVariables;
+      variables?: TVariables;
       errorPolicy: ErrorPolicy;
-      updateQueries: MutationOptions<TData>["updateQueries"],
-      update?: (
-        cache: ApolloCache<TStore>,
-        result: FetchResult<TData>,
-      ) => void;
+      context?: TContext;
+      updateQueries: UpdateQueries<TData>,
+      update?: MutationUpdaterFunction<TData, TVariables, TContext, TCache>;
     },
   ) {
     const data = typeof optimisticResponse === "function"
@@ -409,7 +418,7 @@ export class QueryManager<TStore> {
 
     return this.cache.recordOptimisticTransaction(cache => {
       try {
-        this.markMutationResult<TData>({
+        this.markMutationResult<TData, TVariables, TContext, TCache>({
           ...mutation,
           result: { data },
         }, cache);
@@ -505,9 +514,9 @@ export class QueryManager<TStore> {
     return transformCache.get(document)!;
   }
 
-  private getVariables(
+  private getVariables<TVariables>(
     document: DocumentNode,
-    variables?: OperationVariables,
+    variables?: TVariables,
   ): OperationVariables {
     return {
       ...this.transform(document).defaultVars,

src/core/__tests__/QueryManager/index.ts

@@ -5381,7 +5381,7 @@ describe('QueryManager', () => {
 
   describe('awaitRefetchQueries', () => {
     const awaitRefetchTest =
-    ({ awaitRefetchQueries, testQueryError = false }: MutationBaseOptions & { testQueryError?: boolean }) =>
+    ({ awaitRefetchQueries, testQueryError = false }: MutationBaseOptions<any, any, any> & { testQueryError?: boolean }) =>
     new Promise((resolve, reject) => {
       const query = gql`
         query getAuthors($id: ID!) {
@@ -5455,7 +5455,7 @@ describe('QueryManager', () => {
         { observable },
         result => {
           expect(stripSymbols(result.data)).toEqual(queryData);
-          const mutateOptions: MutationOptions = {
+          const mutateOptions: MutationOptions<any, any, any> = {
             mutation,
             refetchQueries: ['getAuthors'],
           };

src/core/index.ts

@@ -21,7 +21,6 @@ export {
   ErrorPolicy,
   FetchMoreQueryOptions,
   SubscribeToMoreOptions,
-  MutationUpdaterFn,
 } from './watchQueryOptions';
 export { NetworkStatus } from './networkStatus';
 export * from './types';

src/core/types.ts

@@ -1,5 +1,6 @@
 import { DocumentNode, GraphQLError } from 'graphql';
 
+import { ApolloCache } from '../cache';
 import { FetchResult } from '../link/core';
 import { ApolloError } from '../errors';
 import { QueryInfo } from './QueryInfo';
@@ -10,6 +11,8 @@ import { Cache } from '../cache';
 
 export { TypedDocumentNode } from '@graphql-typed-document-node/core';
 
+export type DefaultContext = Record<string, any>;
+
 export type QueryListener = (queryInfo: QueryInfo) => void;
 
 export type ReobserveQueryCallback = (
@@ -51,6 +54,19 @@ export type MutationQueryReducersMap<T = { [key: string]: any }> = {
   [queryName: string]: MutationQueryReducer<T>;
 };
 
+export type MutationUpdaterFunction<
+  TData,
+  TVariables,
+  TContext,
+  TCache extends ApolloCache<any>
+> = (
+  cache: TCache,
+  result: Omit<FetchResult<TData>, 'context'>,
+  options: {
+    context?: TContext,
+    variables?: TVariables,
+  },
+) => void;
 export interface Resolvers {
   [key: string]: {
     [ field: string ]: Resolver;