Refactor markMutationResult to be a method of QueryManager

src/core/QueryManager.ts

@@ -35,19 +35,13 @@ import { NetworkStatus, isNetworkRequestInFlight } from './networkStatus';
 import {
   ApolloQueryResult,
   OperationVariables,
-  MutationQueryReducer,
 } from './types';
 import { LocalState } from './LocalState';
 
 import { QueryInfo, QueryStoreValue, shouldWriteResult } from './QueryInfo';
 
 const { hasOwnProperty } = Object.prototype;
 
-type QueryWithUpdater = {
-  updater: MutationQueryReducer<Object>;
-  queryInfo: QueryInfo;
-};
-
 export class QueryManager<TStore> {
   public cache: ApolloCache<TStore>;
   public link: ApolloLink;
@@ -121,7 +115,7 @@ export class QueryManager<TStore> {
     mutation,
     variables,
     optimisticResponse,
-    updateQueries: updateQueriesByName,
+    updateQueries,
     refetchQueries = [],
     awaitRefetchQueries = false,
     update: updateWithProxyFn,
@@ -148,58 +142,21 @@ export class QueryManager<TStore> {
       variables = await this.localState.addExportedVariables(mutation, variables, context);
     }
 
-    // Create a map of update queries by id to the query instead of by name.
-    const generateUpdateQueriesInfo: () => {
-      [queryId: string]: QueryWithUpdater;
-    } = () => {
-      const ret: { [queryId: string]: QueryWithUpdater } = {};
-
-      if (updateQueriesByName) {
-        this.queries.forEach(({ observableQuery }, queryId) => {
-          if (observableQuery) {
-            const { queryName } = observableQuery;
-            if (
-              queryName &&
-              hasOwnProperty.call(updateQueriesByName, queryName)
-            ) {
-              ret[queryId] = {
-                updater: updateQueriesByName[queryName],
-                queryInfo: this.queries.get(queryId)!,
-              };
-            }
-          }
-        });
-      }
-
-      return ret;
-    };
-
     this.mutationStore.initMutation(
       mutationId,
       mutation,
       variables,
     );
 
     if (optimisticResponse) {
-      const optimistic = typeof optimisticResponse === 'function'
-        ? optimisticResponse(variables)
-        : optimisticResponse;
-
-      this.cache.recordOptimisticTransaction(cache => {
-        try {
-          markMutationResult({
-            mutationId: mutationId,
-            result: { data: optimistic },
-            document: mutation,
-            variables: variables,
-            errorPolicy,
-            queryUpdatersById: generateUpdateQueriesInfo(),
-            update: updateWithProxyFn,
-          }, cache);
-        } catch (error) {
-          invariant.error(error);
-        }
-      }, mutationId);
+      this.markMutationOptimistic<T>(optimisticResponse, {
+        mutationId,
+        document: mutation,
+        variables,
+        errorPolicy,
+        updateQueries,
+        update: updateWithProxyFn,
+      });
     }
 
     this.broadcastQueries();
@@ -231,15 +188,15 @@ export class QueryManager<TStore> {
 
           if (fetchPolicy !== 'no-cache') {
             try {
-              markMutationResult({
+              self.markMutationResult<T>({
                 mutationId,
                 result,
                 document: mutation,
                 variables,
                 errorPolicy,
-                queryUpdatersById: generateUpdateQueriesInfo(),
+                updateQueries,
                 update: updateWithProxyFn,
-              }, self.cache);
+              });
             } catch (e) {
               error = new ApolloError({
                 networkError: e,
@@ -333,6 +290,112 @@ export class QueryManager<TStore> {
     });
   }
 
+  public markMutationResult<TData>(
+    mutation: {
+      mutationId: string;
+      result: FetchResult<TData>;
+      document: DocumentNode;
+      variables?: OperationVariables;
+      errorPolicy: ErrorPolicy;
+      updateQueries: MutationOptions<TData>["updateQueries"],
+      update?: (
+        cache: ApolloCache<TStore>,
+        result: FetchResult<TData>,
+      ) => void;
+    },
+    cache = this.cache,
+  ) {
+    if (shouldWriteResult(mutation.result, mutation.errorPolicy)) {
+      const cacheWrites: Cache.WriteOptions[] = [{
+        result: mutation.result.data,
+        dataId: 'ROOT_MUTATION',
+        query: mutation.document,
+        variables: mutation.variables,
+      }];
+
+      const { updateQueries } = mutation;
+      if (updateQueries) {
+        this.queries.forEach(({ observableQuery }, queryId) => {
+          const queryName = observableQuery && observableQuery.queryName;
+          if (!queryName || !hasOwnProperty.call(updateQueries, queryName)) {
+            return;
+          }
+          const updater = updateQueries[queryName];
+          const { document, variables } = this.queries.get(queryId)!;
+
+          // Read the current query result from the store.
+          const { result: currentQueryResult, complete } = cache.diff<TData>({
+            query: document!,
+            variables,
+            returnPartialData: true,
+            optimistic: false,
+          });
+
+          if (complete && currentQueryResult) {
+            // Run our reducer using the current query result and the mutation result.
+            const nextQueryResult = updater(currentQueryResult, {
+              mutationResult: mutation.result,
+              queryName: document && getOperationName(document) || void 0,
+              queryVariables: variables!,
+            });
+
+            // Write the modified result back into the store if we got a new result.
+            if (nextQueryResult) {
+              cacheWrites.push({
+                result: nextQueryResult,
+                dataId: 'ROOT_QUERY',
+                query: document!,
+                variables,
+              });
+            }
+          }
+        });
+      }
+
+      cache.performTransaction(c => {
+        cacheWrites.forEach(write => c.write(write));
+
+        // If the mutation has some writes associated with it then we need to
+        // apply those writes to the store by running this reducer again with a
+        // write action.
+        const { update } = mutation;
+        if (update) {
+          update(c, mutation.result);
+        }
+      }, /* non-optimistic transaction: */ null);
+    }
+  }
+
+  public markMutationOptimistic<TData>(
+    optimisticResponse: any,
+    mutation: {
+      mutationId: string;
+      document: DocumentNode;
+      variables?: OperationVariables;
+      errorPolicy: ErrorPolicy;
+      updateQueries: MutationOptions<TData>["updateQueries"],
+      update?: (
+        cache: ApolloCache<TStore>,
+        result: FetchResult<TData>,
+      ) => void;
+    },
+  ) {
+    const data = typeof optimisticResponse === "function"
+      ? optimisticResponse(mutation.variables)
+      : optimisticResponse;
+
+    return this.cache.recordOptimisticTransaction(cache => {
+      try {
+        this.markMutationResult<TData>({
+          ...mutation,
+          result: { data },
+        }, cache);
+      } catch (error) {
+        invariant.error(error);
+      }
+    }, mutation.mutationId);
+  }
+
   public fetchQuery<TData, TVars>(
     queryId: string,
     options: WatchQueryOptions<TVars, TData>,
@@ -1075,80 +1138,3 @@ export class QueryManager<TStore> {
     };
   }
 }
-
-function markMutationResult<TStore, TData>(
-  mutation: {
-    mutationId: string;
-    result: FetchResult<TData>;
-    document: DocumentNode;
-    variables: any;
-    errorPolicy: ErrorPolicy;
-    queryUpdatersById: Record<string, QueryWithUpdater>;
-    update:
-      ((cache: ApolloCache<TStore>, mutationResult: Object) => void) |
-      undefined;
-  },
-  cache: ApolloCache<TStore>,
-) {
-  // Incorporate the result from this mutation into the store
-  if (shouldWriteResult(mutation.result, mutation.errorPolicy)) {
-    const cacheWrites: Cache.WriteOptions[] = [{
-      result: mutation.result.data,
-      dataId: 'ROOT_MUTATION',
-      query: mutation.document,
-      variables: mutation.variables,
-    }];
-
-    const { queryUpdatersById } = mutation;
-    if (queryUpdatersById) {
-      Object.keys(queryUpdatersById).forEach(id => {
-        const {
-          updater,
-          queryInfo: {
-            document,
-            variables,
-          },
-        } = queryUpdatersById[id];
-
-        // Read the current query result from the store.
-        const { result: currentQueryResult, complete } = cache.diff<TData>({
-          query: document!,
-          variables,
-          returnPartialData: true,
-          optimistic: false,
-        });
-
-        if (complete && currentQueryResult) {
-          // Run our reducer using the current query result and the mutation result.
-          const nextQueryResult = updater(currentQueryResult, {
-            mutationResult: mutation.result,
-            queryName: getOperationName(document!) || undefined,
-            queryVariables: variables!,
-          });
-
-          // Write the modified result back into the store if we got a new result.
-          if (nextQueryResult) {
-            cacheWrites.push({
-              result: nextQueryResult,
-              dataId: 'ROOT_QUERY',
-              query: document!,
-              variables,
-            });
-          }
-        }
-      });
-    }
-
-    cache.performTransaction(c => {
-      cacheWrites.forEach(write => c.write(write));
-
-      // If the mutation has some writes associated with it then we need to
-      // apply those writes to the store by running this reducer again with a
-      // write action.
-      const { update } = mutation;
-      if (update) {
-        update(c, mutation.result);
-      }
-    }, /* non-optimistic transaction: */ null);
-  }
-}