Use cache.batch within cache.updateQuery and cache.updateFragment

CHANGELOG.md

@@ -36,6 +36,9 @@
 - A new nested entry point called `@apollo/client/testing/core` has been created. Importing from this entry point instead of `@apollo/client/testing` excludes any React-related dependencies. <br/>
   [@wassim-k](https://github.com/wassim-k) in [#8687](https://github.com/apollographql/apollo-client/pull/8687)
 
+- Make `cache.batch` return the result of calling the `options.update` function. <br/>
+  [@benjamn](https://github.com/benjamn) in [#8696](https://github.com/apollographql/apollo-client/pull/8696)
+
 ### React Refactoring
 
 #### Bug Fixes (due to [@brainkim](https://github.com/brainkim) in [#8596](https://github.com/apollographql/apollo-client/pull/8596)):

src/cache/core/__tests__/cache.ts

@@ -21,6 +21,7 @@ class TestCache extends ApolloCache<unknown> {
   }
 
   public performTransaction(transaction: <TSerialized>(c: ApolloCache<TSerialized>) => void): void {
+    transaction(this);
   }
 
   public read<T, TVariables = any>(query: Cache.ReadOptions<TVariables>): T | null {

src/cache/core/cache.ts

@@ -14,14 +14,16 @@ export type Transaction<T> = (c: ApolloCache<T>) => void;
 export abstract class ApolloCache<TSerialized> implements DataProxy {
   // required to implement
   // core API
-  public abstract read<T, TVariables = any>(
-    query: Cache.ReadOptions<TVariables, T>,
-  ): T | null;
-  public abstract write<TResult = any, TVariables = any>(
-    write: Cache.WriteOptions<TResult, TVariables>,
+  public abstract read<TData = any, TVariables = any>(
+    query: Cache.ReadOptions<TVariables, TData>,
+  ): TData | null;
+  public abstract write<TData = any, TVariables = any>(
+    write: Cache.WriteOptions<TData, TVariables>,
   ): Reference | undefined;
   public abstract diff<T>(query: Cache.DiffOptions): Cache.DiffResult<T>;
-  public abstract watch(watch: Cache.WatchOptions): () => void;
+  public abstract watch<TData = any, TVariables = any>(
+    watch: Cache.WatchOptions<TData, TVariables>,
+  ): () => void;
   public abstract reset(): Promise<void>;
 
   // Remove whole objects from the cache by passing just options.id, or
@@ -59,11 +61,16 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
   // provide a default batch implementation that's just another way of calling
   // performTransaction. Subclasses of ApolloCache (such as InMemoryCache) can
   // override the batch method to do more interesting things with its options.
-  public batch(options: Cache.BatchOptions<this>) {
+  public batch<U>(options: Cache.BatchOptions<this, U>): U {
     const optimisticId =
       typeof options.optimistic === "string" ? options.optimistic :
       options.optimistic === false ? null : void 0;
-    this.performTransaction(options.update, optimisticId);
+    let updateResult: U;
+    this.performTransaction(
+      () => updateResult = options.update(this),
+      optimisticId,
+    );
+    return updateResult!;
   }
 
   public abstract performTransaction(
@@ -171,21 +178,29 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
     options: Cache.UpdateQueryOptions<TData, TVariables>,
     update: (data: TData | null) => TData | null | void,
   ): TData | null {
-    const value = this.readQuery<TData, TVariables>(options);
-    const data = update(value);
-    if (data === void 0 || data === null) return value;
-    this.writeQuery<TData, TVariables>({ ...options, data });
-    return data;
+    return this.batch({
+      update(cache) {
+        const value = cache.readQuery<TData, TVariables>(options);
+        const data = update(value);
+        if (data === void 0 || data === null) return value;
+        cache.writeQuery<TData, TVariables>({ ...options, data });
+        return data;
+      },
+    });
   }
 
   public updateFragment<TData = any, TVariables = any>(
     options: Cache.UpdateFragmentOptions<TData, TVariables>,
     update: (data: TData | null) => TData | null | void,
   ): TData | null {
-    const value = this.readFragment<TData, TVariables>(options);
-    const data = update(value);
-    if (data === void 0 || data === null) return value;
-    this.writeFragment<TData, TVariables>({ ...options, data });
-    return data;
+    return this.batch({
+      update(cache) {
+        const value = cache.readFragment<TData, TVariables>(options);
+        const data = update(value);
+        if (data === void 0 || data === null) return value;
+        cache.writeFragment<TData, TVariables>({ ...options, data });
+        return data;
+      },
+    });
   }
 }

src/cache/core/types/Cache.ts

@@ -3,9 +3,9 @@ import { Modifier, Modifiers } from './common';
 import { ApolloCache } from '../cache';
 
 export namespace Cache {
-  export type WatchCallback = (
-    diff: Cache.DiffResult<any>,
-    lastDiff?: Cache.DiffResult<any>,
+  export type WatchCallback<TData = any> = (
+    diff: Cache.DiffResult<TData>,
+    lastDiff?: Cache.DiffResult<TData>,
   ) => void;
 
   export interface ReadOptions<TVariables = any, TData = any>
@@ -25,19 +25,23 @@ export namespace Cache {
     result: TResult;
   }
 
-  export interface DiffOptions extends ReadOptions {
+  export interface DiffOptions<
+    TData = any,
+    TVariables = any,
+  > extends ReadOptions<TVariables, TData> {
     // The DiffOptions interface is currently just an alias for
     // ReadOptions, though DiffOptions used to be responsible for
     // declaring the returnPartialData option.
   }
 
   export interface WatchOptions<
-    Watcher extends object = Record<string, any>
-  > extends ReadOptions {
-    watcher?: Watcher;
+    TData = any,
+    TVariables = any,
+  > extends ReadOptions<TVariables, TData> {
+    watcher?: object;
     immediate?: boolean;
-    callback: WatchCallback;
-    lastDiff?: DiffResult<any>;
+    callback: WatchCallback<TData>;
+    lastDiff?: DiffResult<TData>;
   }
 
   export interface EvictOptions {
@@ -54,10 +58,10 @@ export namespace Cache {
     broadcast?: boolean;
   }
 
-  export interface BatchOptions<C extends ApolloCache<any>> {
+  export interface BatchOptions<C extends ApolloCache<any>, U = void> {
     // Same as the first parameter of performTransaction, except the cache
     // argument will have the subclass type rather than ApolloCache.
-    update(cache: C): void;
+    update(cache: C): U;
 
     // Passing a string for this option creates a new optimistic layer, with the
     // given string as its layer.id, just like passing a string for the
@@ -66,7 +70,7 @@ export namespace Cache {
     // against the current top layer of the cache), and passing false is the
     // same as passing null (running the operation against root/non-optimistic
     // cache data).
-    optimistic: string | boolean;
+    optimistic?: string | boolean;
 
     // If you specify the ID of an optimistic layer using this option, that
     // layer will be removed as part of the batch transaction, triggering at