From bf0047255cd83c611b4f2054243eab1fa4a80a84 Mon Sep 17 00:00:00 2001
From: Swain Molster <swain.molster@lifeomic.com>
Date: Wed, 11 Dec 2024 11:17:47 -0500
Subject: [PATCH] feat: allow updating cache entries using a predicate function

---
 src/cache.ts       | 56 ++++++++++++++++++++++++++++++----------------
 src/hooks.test.tsx | 44 ++++++++++++++++++++++++++++++++++++
 src/hooks.ts       |  4 +---
 src/types.ts       |  8 +++++--
 4 files changed, 88 insertions(+), 24 deletions(-)

diff --git a/src/cache.ts b/src/cache.ts
index 951fff5..6aa53e2 100644
--- a/src/cache.ts
+++ b/src/cache.ts
@@ -2,13 +2,8 @@ import { QueryClient, QueryFilters } from '@tanstack/react-query';
 // eslint-disable-next-line no-restricted-imports
 import { isEqual } from 'lodash';
 import { produce } from 'immer';
-import {
-  CacheUtils,
-  EndpointInvalidationMap,
-  RequestPayloadOf,
-  RoughEndpoints,
-} from './types';
-import { InternalQueryKey, isInternalQueryKey } from './util';
+import { CacheUtils, EndpointInvalidationMap, RoughEndpoints } from './types';
+import { createQueryKey, isInternalQueryKey } from './util';
 
 const createQueryFilterFromSpec = <Endpoints extends RoughEndpoints>(
   endpoints: EndpointInvalidationMap<Endpoints>,
@@ -47,23 +42,43 @@ export const INFINITE_QUERY_KEY = 'infinite' as const;
 
 export const createCacheUtils = <Endpoints extends RoughEndpoints>(
   client: QueryClient,
-  makeQueryKey: <Route extends keyof Endpoints & string>(
-    route: Route,
-    payload: RequestPayloadOf<Endpoints, Route>,
-  ) => InternalQueryKey,
+  name: string,
 ): CacheUtils<Endpoints> => {
   const updateCache: (
     keyPrefix?: typeof INFINITE_QUERY_KEY,
   ) => CacheUtils<Endpoints>['updateCache'] =
-    (keyPrefix) => (route, payload, updater) => {
-      client.setQueryData<Endpoints[typeof route]['Response']>(
-        [keyPrefix, makeQueryKey(route, payload)].filter(Boolean),
+    (keyPrefix) => (route, payloadOrPredicate, updater) => {
+      client.setQueriesData<Endpoints[typeof route]['Response']>(
+        typeof payloadOrPredicate === 'function'
+          ? {
+              predicate: ({ queryKey }) => {
+                /* istanbul ignore next */
+                if (keyPrefix && queryKey[0] !== keyPrefix) {
+                  /* istanbul ignore next */
+                  return false;
+                }
+                const payloadInKey = keyPrefix ? queryKey[1] : queryKey[0];
+
+                return (
+                  isInternalQueryKey(payloadInKey) &&
+                  payloadInKey.name === name &&
+                  payloadInKey.route === route &&
+                  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                  // @ts-ignore These types are a bit finicky. Override it.
+                  payloadOrPredicate(payloadInKey.payload)
+                );
+              },
+            }
+          : {
+              queryKey: [
+                keyPrefix,
+                createQueryKey(name, route, payloadOrPredicate),
+              ].filter(Boolean),
+              exact: true,
+            },
         typeof updater !== 'function'
           ? updater
           : (current) => {
-              if (current === undefined) {
-                return;
-              }
               return produce(
                 current,
                 // @ts-expect-error TypeScript incorrectly thinks that `updater`
@@ -84,9 +99,12 @@ export const createCacheUtils = <Endpoints extends RoughEndpoints>(
     updateCache: updateCache(),
     updateInfiniteCache: updateCache(INFINITE_QUERY_KEY),
     getQueryData: (route, payload) =>
-      client.getQueryData([makeQueryKey(route, payload)]),
+      client.getQueryData([createQueryKey(name, route, payload)]),
     getInfiniteQueryData: (route, payload) =>
-      client.getQueryData([INFINITE_QUERY_KEY, makeQueryKey(route, payload)]),
+      client.getQueryData([
+        INFINITE_QUERY_KEY,
+        createQueryKey(name, route, payload),
+      ]),
     getQueriesData: (route) =>
       client
         .getQueriesData(createQueryFilterFromSpec({ [route]: 'all' }))
diff --git a/src/hooks.test.tsx b/src/hooks.test.tsx
index 2064e3d..a119687 100644
--- a/src/hooks.test.tsx
+++ b/src/hooks.test.tsx
@@ -1557,6 +1557,50 @@ describe('useAPICache', () => {
 
         expect(updateFn).not.toHaveBeenCalled();
       });
+
+      it('allows updating using a predicate', async () => {
+        const updater =
+          method === 'updateCache'
+            ? () => ({ message: 'Samwise Gamgee' })
+            : () => ({
+                pages: [{ items: [{ message: 'Samwise Gamgee' }] }],
+                pageParams: [],
+              });
+
+        const screen = render(() => (
+          <TestComponent
+            getRenderData={config.getRenderData}
+            onPress={(cache) => {
+              const updateMethod = cache[method];
+
+              updateMethod(
+                // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                // @ts-expect-error
+                config.route,
+                // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                // @ts-expect-error
+                (payload) => payload.filter === '',
+                updater,
+              );
+            }}
+          />
+        ));
+
+        await screen.findByText('Response: Frodo Baggins');
+
+        expect(client.request).toHaveBeenCalledTimes(1);
+
+        TestingLibrary.fireEvent.click(screen.getByText('Update Cache'));
+
+        // The update does not happen immediately.
+        await TestingLibrary.waitFor(() => {
+          expect(screen.getByTestId('render-data').textContent).toStrictEqual(
+            'Response: Samwise Gamgee',
+          );
+        });
+
+        expect(client.request).toHaveBeenCalledTimes(1);
+      });
     });
   });
 });
diff --git a/src/hooks.ts b/src/hooks.ts
index 34ec239..e1437e1 100644
--- a/src/hooks.ts
+++ b/src/hooks.ts
@@ -181,9 +181,7 @@ export const createAPIHooks = <Endpoints extends RoughEndpoints>({
 
     useAPICache: () => {
       const client = useQueryClient();
-      return createCacheUtils(client, (route, payload) =>
-        createQueryKey(name, route, payload),
-      );
+      return createCacheUtils(client, name);
     },
   };
 };
diff --git a/src/types.ts b/src/types.ts
index 34890d3..d0d890e 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -140,13 +140,17 @@ export type CacheUtils<Endpoints extends RoughEndpoints> = {
 
   updateCache: <Route extends keyof Endpoints & string>(
     route: Route,
-    payload: RequestPayloadOf<Endpoints, Route>,
+    payloadOrPredicate:
+      | RequestPayloadOf<Endpoints, Route>
+      | ((payload: RequestPayloadOf<Endpoints, Route>) => boolean),
     updater: CacheUpdate<Endpoints[Route]['Response']>,
   ) => void;
 
   updateInfiniteCache: <Route extends keyof Endpoints & string>(
     route: Route,
-    payload: RequestPayloadOf<Endpoints, Route>,
+    payloadOrPredicate:
+      | RequestPayloadOf<Endpoints, Route>
+      | ((payload: RequestPayloadOf<Endpoints, Route>) => boolean),
     updater: CacheUpdate<InfiniteData<Endpoints[Route]['Response']>>,
   ) => void;