feat: allow updating cache entries using a predicate function

lifeomic · Dec 11, 2024 · 4b99c20 · 4b99c20
1 parent 520cf2f
commit 4b99c20
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 20 deletions.
diff --git 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,17 +42,44 @@ 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 }) => {
+                if (keyPrefix && queryKey[0] !== keyPrefix) {
+                  return false;
+                }
+
+                const payloadInKey = keyPrefix ? queryKey[1] : queryKey[0];
+
+                if (!isInternalQueryKey(payloadInKey)) {
+                  return false;
+                }
+
+                if (payloadInKey.name !== name) {
+                  return false;
+                }
+
+                if (payloadInKey.route !== route) {
+                  return false;
+                }
+
+                return payloadOrPredicate(payloadInKey.payload as any);
+              },
+            }
+          : {
+              queryKey: [
+                keyPrefix,
+                createQueryKey(name, route, payloadOrPredicate),
+              ].filter(Boolean),
+              exact: true,
+            },
         typeof updater !== 'function'
           ? updater
           : (current) => {
@@ -84,9 +106,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.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
@@ -140,7 +140,9 @@ 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;