fix(experimental_createQueryPersister): Change persisterRestoreAll to restoreQueries with filter support (#9230)

* Add persisterRestoreByKey

* Fix imports

* Add missing return statement

* Fix test

* docs

* Add extra test

* Merge into single restoreQueries function

Author: rklomp

docs/framework/react/plugins/createPersister.md

@@ -111,10 +111,17 @@ This function can be used to sporadically clean up stoage from `expired`, `buste
 For this function to work, your storage must expose `entries` method that would return a `key-value tuple array`.  
 For example `Object.entries(localStorage)` for `localStorage` or `entries` from `idb-keyval`.
 
-### `persisterRestoreAll(queryClient: QueryClient): Promise<void>`
+### `restoreQueries(queryClient: QueryClient, filters): Promise<void>`
 
-This function can be used to restore all queries that are currently stored by persister in one go.  
-For example when your app is starting up in offline mode, or you want data from previous session to be immediately available without intermediate `loading` state.
+This function can be used to restore queries that are currently stored by persister.  
+For example when your app is starting up in offline mode, or you want all or only specific data from previous session to be immediately available without intermediate `loading` state.
+
+The filter object supports the following properties:
+
+- `queryKey?: QueryKey`
+  - Set this property to define a query key to match on.
+- `exact?: boolean`
+  - If you don't want to search queries inclusively by query key, you can pass the `exact: true` option to return only the query with the exact query key you have passed.
 
 For this function to work, your storage must expose `entries` method that would return a `key-value tuple array`.  
 For example `Object.entries(localStorage)` for `localStorage` or `entries` from `idb-keyval`.

packages/query-core/src/index.ts

@@ -15,6 +15,7 @@ export { focusManager } from './focusManager'
 export { onlineManager } from './onlineManager'
 export {
   hashKey,
+  partialMatchKey,
   replaceEqualDeep,
   isServer,
   matchQuery,

packages/query-persist-client-core/src/__tests__/createPersister.test.ts

@@ -497,7 +497,7 @@ describe('createPersister', () => {
     })
   })
 
-  describe('persisterRestoreAll', () => {
+  describe('restoreQueries', () => {
     test('should properly clean storage from busted entries', async () => {
       const storage = getFreshStorage()
       const { persister, client, query, queryKey } = setupPersister(['foo'], {
@@ -513,11 +513,11 @@ describe('createPersister', () => {
 
       expect(await storage.entries()).toHaveLength(1)
 
-      await persister.persisterRestoreAll(client)
+      await persister.restoreQueries(client)
       expect(await storage.entries()).toHaveLength(0)
     })
 
-    test('should properly restore queries from cache', async () => {
+    test('should properly restore queries from cache without filters', async () => {
       const storage = getFreshStorage()
       const { persister, client, queryKey } = setupPersister(['foo'], {
         storage,
@@ -530,10 +530,102 @@ describe('createPersister', () => {
       client.clear()
       expect(client.getQueryCache().getAll()).toHaveLength(0)
 
-      await persister.persisterRestoreAll(client)
+      await persister.restoreQueries(client)
+      expect(client.getQueryCache().getAll()).toHaveLength(1)
+
+      expect(client.getQueryData(queryKey)).toEqual('foo')
+    })
+
+    test('should properly restore queries from cache', async () => {
+      const storage = getFreshStorage()
+      const { persister, client, queryKey } = setupPersister(['foo', 'bar'], {
+        storage,
+      })
+      client.setQueryData(queryKey, 'foo')
+
+      await persister.persistQueryByKey(queryKey, client)
+
+      expect(await storage.entries()).toHaveLength(1)
+      client.clear()
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+
+      await persister.restoreQueries(client, { queryKey })
       expect(client.getQueryCache().getAll()).toHaveLength(1)
 
       expect(client.getQueryData(queryKey)).toEqual('foo')
     })
+
+    test('should not restore queries from cache if there is no match', async () => {
+      const storage = getFreshStorage()
+      const { persister, client, queryKey } = setupPersister(['foo', 'bar'], {
+        storage,
+      })
+      client.setQueryData(queryKey, 'foo')
+
+      await persister.persistQueryByKey(queryKey, client)
+
+      expect(await storage.entries()).toHaveLength(1)
+      client.clear()
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+
+      await persister.restoreQueries(client, { queryKey: ['bar'] })
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+    })
+
+    test('should properly restore queries from cache with partial match', async () => {
+      const storage = getFreshStorage()
+      const { persister, client, queryKey } = setupPersister(['foo', 'bar'], {
+        storage,
+      })
+      client.setQueryData(queryKey, 'foo')
+
+      await persister.persistQueryByKey(queryKey, client)
+
+      expect(await storage.entries()).toHaveLength(1)
+      client.clear()
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+
+      await persister.restoreQueries(client, { queryKey: ['foo'] })
+      expect(client.getQueryCache().getAll()).toHaveLength(1)
+
+      expect(client.getQueryData(queryKey)).toEqual('foo')
+    })
+
+    test('should not restore queries from cache with exact match if there is no match', async () => {
+      const storage = getFreshStorage()
+      const { persister, client, queryKey } = setupPersister(['foo', 'bar'], {
+        storage,
+      })
+      client.setQueryData(queryKey, 'foo')
+
+      await persister.persistQueryByKey(queryKey, client)
+
+      expect(await storage.entries()).toHaveLength(1)
+      client.clear()
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+
+      await persister.restoreQueries(client, { queryKey: ['foo'], exact: true })
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+    })
+
+    test('should restore queries from cache with exact match', async () => {
+      const storage = getFreshStorage()
+      const { persister, client, queryKey } = setupPersister(['foo', 'bar'], {
+        storage,
+      })
+      client.setQueryData(queryKey, 'foo')
+
+      await persister.persistQueryByKey(queryKey, client)
+
+      expect(await storage.entries()).toHaveLength(1)
+      client.clear()
+      expect(client.getQueryCache().getAll()).toHaveLength(0)
+
+      await persister.restoreQueries(client, {
+        queryKey: queryKey,
+        exact: true,
+      })
+      expect(client.getQueryCache().getAll()).toHaveLength(1)
+    })
   })
 })

packages/query-persist-client-core/src/createPersister.ts

@@ -1,4 +1,4 @@
-import { matchQuery } from '@tanstack/query-core'
+import { hashKey, matchQuery, partialMatchKey } from '@tanstack/query-core'
 import type {
   Query,
   QueryClient,
@@ -240,7 +240,12 @@ export function experimental_createQueryPersister<TStorageValue = string>({
     }
   }
 
-  async function persisterRestoreAll(queryClient: QueryClient) {
+  async function restoreQueries(
+    queryClient: QueryClient,
+    filters: Pick<QueryFilters, 'queryKey' | 'exact'> = {},
+  ): Promise<void> {
+    const { exact, queryKey } = filters
+
     if (storage?.entries) {
       const entries = await storage.entries()
       for (const [key, value] of entries) {
@@ -249,15 +254,26 @@ export function experimental_createQueryPersister<TStorageValue = string>({
 
           if (isExpiredOrBusted(persistedQuery)) {
             await storage.removeItem(key)
-          } else {
-            queryClient.setQueryData(
-              persistedQuery.queryKey,
-              persistedQuery.state.data,
-              {
-                updatedAt: persistedQuery.state.dataUpdatedAt,
-              },
-            )
+            continue
+          }
+
+          if (queryKey) {
+            if (exact) {
+              if (persistedQuery.queryHash !== hashKey(queryKey)) {
+                continue
+              }
+            } else if (!partialMatchKey(persistedQuery.queryKey, queryKey)) {
+              continue
+            }
           }
+
+          queryClient.setQueryData(
+            persistedQuery.queryKey,
+            persistedQuery.state.data,
+            {
+              updatedAt: persistedQuery.state.dataUpdatedAt,
+            },
+          )
         }
       }
     } else if (process.env.NODE_ENV === 'development') {
@@ -273,6 +289,6 @@ export function experimental_createQueryPersister<TStorageValue = string>({
     persistQueryByKey,
     retrieveQuery,
     persisterGc,
-    persisterRestoreAll,
+    restoreQueries,
   }
 }