From a4d4574f5d18476e0a468c00ec8e7dbbae8ca087 Mon Sep 17 00:00:00 2001
From: wassim-k <wassim.khalil9@gmail.com>
Date: Sun, 13 Jun 2021 16:12:48 +1000
Subject: [PATCH] Add modifyQuery & modifyFragment to ApolloCache

---
 src/cache/core/__tests__/cache.ts | 163 +++++++++++++++++++++++++++++-
 src/cache/core/cache.ts           |  22 ++++
 src/cache/core/types/Cache.ts     |   2 +
 src/cache/core/types/DataProxy.ts |   6 ++
 4 files changed, 189 insertions(+), 4 deletions(-)

diff --git a/src/cache/core/__tests__/cache.ts b/src/cache/core/__tests__/cache.ts
index f1417add255..5098216ef56 100644
--- a/src/cache/core/__tests__/cache.ts
+++ b/src/cache/core/__tests__/cache.ts
@@ -1,5 +1,5 @@
 import gql from 'graphql-tag';
-import { ApolloCache  } from '../cache';
+import { ApolloCache } from '../cache';
 import { Cache, DataProxy } from '../..';
 import { Reference } from '../../../utilities/graphql/storeUtils';
 
@@ -73,7 +73,7 @@ describe('abstract cache', () => {
       const test = new TestCache();
       test.read = jest.fn();
 
-      test.readQuery({query});
+      test.readQuery({ query });
       expect(test.read).toBeCalled();
     });
 
@@ -81,8 +81,8 @@ describe('abstract cache', () => {
       const test = new TestCache();
       test.read = ({ optimistic }) => optimistic as any;
 
-      expect(test.readQuery({query})).toBe(false);
-      expect(test.readQuery({query}, true)).toBe(true);
+      expect(test.readQuery({ query })).toBe(false);
+      expect(test.readQuery({ query }, true)).toBe(true);
     });
   });
 
@@ -151,4 +151,159 @@ describe('abstract cache', () => {
       expect(test.write).toBeCalled();
     });
   });
+
+  describe('modifyQuery', () => {
+    it('runs the readQuery & writeQuery methods', () => {
+      const test = new TestCache();
+      test.readQuery = jest.fn();
+      test.writeQuery = jest.fn();
+
+      test.modifyQuery({ query }, data => 'foo');
+
+      expect(test.readQuery).toBeCalled();
+      expect(test.writeQuery).toBeCalled();
+    });
+
+    it('does not call writeQuery method if data is null', () => {
+      const test = new TestCache();
+      test.readQuery = jest.fn();
+      test.writeQuery = jest.fn();
+
+      test.modifyQuery({ query }, data => null);
+
+      expect(test.readQuery).toBeCalled();
+      expect(test.writeQuery).not.toBeCalled();
+    });
+
+    it('does not call writeQuery method if data is undefined', () => {
+      const test = new TestCache();
+      test.readQuery = jest.fn();
+      test.writeQuery = jest.fn();
+
+      test.modifyQuery({ query }, data => { return; });
+
+      expect(test.readQuery).toBeCalled();
+      expect(test.writeQuery).not.toBeCalled();
+    });
+
+    it('calls the readQuery & writeQuery methods with the options object', () => {
+      const test = new TestCache();
+      const options: Cache.ModifyQueryOptions<string, any> = { query, broadcast: true, variables: { test: 1 }, optimistic: true, returnPartialData: true };
+      test.readQuery = jest.fn();
+      test.writeQuery = jest.fn();
+
+      test.modifyQuery(options, data => 'foo');
+
+      expect(test.readQuery).toBeCalledWith(
+        expect.objectContaining(options)
+      );
+
+      expect(test.writeQuery).toBeCalledWith(
+        expect.objectContaining({ ...options, data: 'foo' })
+      );
+    });
+
+    it('returns current value in memory if no update was made', () => {
+      const test = new TestCache();
+      test.readQuery = jest.fn().mockReturnValue('foo');
+      expect(test.modifyQuery({ query }, data => null)).toBe('foo');
+    });
+
+    it('returns the updated value in memory if an update was made', () => {
+      const test = new TestCache();
+      let currentValue = 'foo';
+      test.readQuery = jest.fn().mockImplementation(() => currentValue);
+      test.writeQuery = jest.fn().mockImplementation(({ data }) => currentValue = data);
+      expect(test.modifyQuery({ query }, data => 'bar')).toBe('bar');
+    });
+
+    it('calls modify function with the current value in memory', () => {
+      const test = new TestCache();
+      test.readQuery = jest.fn().mockReturnValue('foo');
+      test.modifyQuery({ query }, data => {
+        expect(data).toBe('foo');
+      });
+    });
+  });
+
+  describe('modifyFragment', () => {
+    const fragmentId = 'frag';
+    const fragment = gql`
+      fragment a on b {
+        name
+      }
+    `;
+
+    it('runs the readFragment & writeFragment methods', () => {
+      const test = new TestCache();
+      test.readFragment = jest.fn();
+      test.writeFragment = jest.fn();
+
+      test.modifyFragment({ id: fragmentId, fragment }, data => 'foo');
+
+      expect(test.readFragment).toBeCalled();
+      expect(test.writeFragment).toBeCalled();
+    });
+
+    it('does not call writeFragment method if data is null', () => {
+      const test = new TestCache();
+      test.readFragment = jest.fn();
+      test.writeFragment = jest.fn();
+
+      test.modifyFragment({ id: fragmentId, fragment }, data => null);
+
+      expect(test.readFragment).toBeCalled();
+      expect(test.writeFragment).not.toBeCalled();
+    });
+
+    it('does not call writeFragment method if data is undefined', () => {
+      const test = new TestCache();
+      test.readFragment = jest.fn();
+      test.writeFragment = jest.fn();
+
+      test.modifyFragment({ id: fragmentId, fragment }, data => { return; });
+
+      expect(test.readFragment).toBeCalled();
+      expect(test.writeFragment).not.toBeCalled();
+    });
+
+    it('calls the readFragment & writeFragment methods with the options object', () => {
+      const test = new TestCache();
+      const options: Cache.ModifyFragmentOptions<string, any> = { id: fragmentId, fragment, fragmentName: 'a', broadcast: true, variables: { test: 1 }, optimistic: true, returnPartialData: true };
+      test.readFragment = jest.fn();
+      test.writeFragment = jest.fn();
+
+      test.modifyFragment(options, data => 'foo');
+
+      expect(test.readFragment).toBeCalledWith(
+        expect.objectContaining(options)
+      );
+
+      expect(test.writeFragment).toBeCalledWith(
+        expect.objectContaining({ ...options, data: 'foo' })
+      );
+    });
+
+    it('returns current value in memory if no update was made', () => {
+      const test = new TestCache();
+      test.readFragment = jest.fn().mockReturnValue('foo');
+      expect(test.modifyFragment({ id: fragmentId, fragment }, data => { return; })).toBe('foo');
+    });
+
+    it('returns the updated value in memory if an update was made', () => {
+      const test = new TestCache();
+      let currentValue = 'foo';
+      test.readFragment = jest.fn().mockImplementation(() => currentValue);
+      test.writeFragment = jest.fn().mockImplementation(({ data }) => currentValue = data);
+      expect(test.modifyFragment({ id: fragmentId, fragment }, data => 'bar')).toBe('bar');
+    });
+
+    it('calls modify function with the current value in memory', () => {
+      const test = new TestCache();
+      test.readFragment = jest.fn().mockReturnValue('foo');
+      test.modifyFragment({ id: fragmentId, fragment }, data => {
+        expect(data).toBe('foo');
+      });
+    });
+  });
 });
diff --git a/src/cache/core/cache.ts b/src/cache/core/cache.ts
index db5c63879a8..b90d5718c21 100644
--- a/src/cache/core/cache.ts
+++ b/src/cache/core/cache.ts
@@ -166,4 +166,26 @@ export abstract class ApolloCache<TSerialized> implements DataProxy {
       result: data,
     }));
   }
+
+  public modifyQuery<TData = any, TVariables = any>(
+    options: Cache.ModifyQueryOptions<TData, TVariables>,
+    modifyFn: (data: TData | null) => TData | null | void
+  ): TData | null {
+    const value = this.readQuery<TData, TVariables>(options);
+    const data = modifyFn(value);
+    if (data === undefined || data === null) return value;
+    this.writeQuery<TData, TVariables>({ ...options, data });
+    return data;
+  }
+
+  public modifyFragment<TData = any, TVariables = any>(
+    options: Cache.ModifyFragmentOptions<TData, TVariables>,
+    modifyFn: (data: TData | null) => TData | null | void
+  ): TData | null {
+    const value = this.readFragment<TData, TVariables>(options);
+    const data = modifyFn(value);
+    if (data === undefined || data === null) return value;
+    this.writeFragment<TData, TVariables>({ ...options, data });
+    return data;
+  }
 }
diff --git a/src/cache/core/types/Cache.ts b/src/cache/core/types/Cache.ts
index 98a21afd4f5..d67a8c691eb 100644
--- a/src/cache/core/types/Cache.ts
+++ b/src/cache/core/types/Cache.ts
@@ -92,5 +92,7 @@ export namespace Cache {
   export import ReadFragmentOptions = DataProxy.ReadFragmentOptions;
   export import WriteQueryOptions = DataProxy.WriteQueryOptions;
   export import WriteFragmentOptions = DataProxy.WriteFragmentOptions;
+  export import ModifyQueryOptions = DataProxy.ModifyQueryOptions;
+  export import ModifyFragmentOptions = DataProxy.ModifyFragmentOptions;
   export import Fragment = DataProxy.Fragment;
 }
diff --git a/src/cache/core/types/DataProxy.ts b/src/cache/core/types/DataProxy.ts
index 8e04d01a5df..ede97d488a3 100644
--- a/src/cache/core/types/DataProxy.ts
+++ b/src/cache/core/types/DataProxy.ts
@@ -118,6 +118,12 @@ export namespace DataProxy {
   export interface WriteFragmentOptions<TData, TVariables>
     extends Fragment<TVariables, TData>, WriteOptions<TData> {}
 
+  export interface ModifyQueryOptions<TData, TVariables>
+    extends Omit<ReadQueryOptions<TData, TVariables> & WriteQueryOptions<TData, TVariables>, 'data'> {}
+
+  export interface ModifyFragmentOptions<TData, TVariables>
+    extends Omit<ReadFragmentOptions<TData, TVariables> & WriteFragmentOptions<TData, TVariables>, 'data'> {}
+
   export type DiffResult<T> = {
     result?: T;
     complete?: boolean;