From 33ff8317a9871fbac56913e746753803264f9220 Mon Sep 17 00:00:00 2001
From: Damian Osipiuk <osipiukd+git@gmail.com>
Date: Thu, 25 Jan 2024 23:34:42 +0100
Subject: [PATCH] fix(experimental_createQueryPersister): return more
 utilities, rename persister

---
 .../react/plugins/createPersister.md          |  22 ++-
 docs/framework/vue/plugins/createPersister.md |  22 ++-
 examples/vue/persister/src/Post.vue           |   6 +-
 examples/vue/persister/src/Posts.vue          |   6 +-
 .../src/__tests__/createPersister.test.ts     |  14 +-
 .../src/createPersister.ts                    | 178 +++++++++++++-----
 .../__tests__/fine-grained-persister.test.tsx |  26 +--
 7 files changed, 185 insertions(+), 89 deletions(-)

diff --git a/docs/framework/react/plugins/createPersister.md b/docs/framework/react/plugins/createPersister.md
index 24206b8bfc..c1707f82fa 100644
--- a/docs/framework/react/plugins/createPersister.md
+++ b/docs/framework/react/plugins/createPersister.md
@@ -1,6 +1,6 @@
 ---
 id: createPersister
-title: experimental_createPersister
+title: experimental_createQueryPersister
 ---
 
 ## Installation
@@ -33,8 +33,8 @@ bun add @tanstack/query-persist-client-core
 
 ## Usage
 
-- Import the `experimental_createPersister` function
-- Create a new `experimental_createPersister`
+- Import the `experimental_createQueryPersister` function
+- Create a new `experimental_createQueryPersister`
   - you can pass any `storage` to it that adheres to the `AsyncStorage` or `Storage` interface - the example below uses the async-storage from React Native.
 - Pass that `persister` as an option to your Query. This can be done either by passing it to the `defaultOptions` of the `QueryClient` or to any `useQuery` hook instance.
   - If you pass this `persister` as `defaultOptions`, all queries will be persisted to the provided `storage`. You can additionally narrow this down by passing `filters`. In contrast to the `persistClient` plugin, this will not persist the whole query client as a single item, but each query separately. As a key, the query hash is used.
@@ -47,16 +47,18 @@ Garbage collecting a Query from memory **does not** affect the persisted data. T
 ```tsx
 import AsyncStorage from '@react-native-async-storage/async-storage'
 import { QueryClient } from '@tanstack/react-query'
-import { experimental_createPersister } from '@tanstack/query-persist-client-core'
+import { experimental_createQueryPersister } from '@tanstack/query-persist-client-core'
+
+const persister = experimental_createQueryPersister({
+  storage: AsyncStorage,
+  maxAge: 1000 * 60 * 60 * 12, // 12 hours
+})
 
 const queryClient = new QueryClient({
   defaultOptions: {
     queries: {
       gcTime: 1000 * 30, // 30 seconds
-      persister: experimental_createPersister({
-        storage: AsyncStorage,
-        maxAge: 1000 * 60 * 60 * 12, // 12 hours
-      }),
+      persister: persister.persisterFn,
     },
   },
 })
@@ -68,10 +70,10 @@ The `createPersister` plugin technically wraps the `queryFn`, so it doesn't rest
 
 ## API
 
-### `experimental_createPersister`
+### `experimental_createQueryPersister`
 
 ```tsx
-experimental_createPersister(options: StoragePersisterOptions)
+experimental_createQueryPersister(options: StoragePersisterOptions)
 ```
 
 #### `Options`
diff --git a/docs/framework/vue/plugins/createPersister.md b/docs/framework/vue/plugins/createPersister.md
index b2f9773ca8..02fb5f0868 100644
--- a/docs/framework/vue/plugins/createPersister.md
+++ b/docs/framework/vue/plugins/createPersister.md
@@ -1,6 +1,6 @@
 ---
 id: createPersister
-title: experimental_createPersister
+title: experimental_createQueryPersister
 ---
 
 ## Installation
@@ -31,8 +31,8 @@ bun add @tanstack/query-persist-client-core
 
 ## Usage
 
-- Import the `experimental_createPersister` function
-- Create a new `experimental_createPersister`
+- Import the `experimental_createQueryPersister` function
+- Create a new `experimental_createQueryPersister`
   - you can pass any `storage` to it that adheres to the `AsyncStorage` or `Storage` interface
 - Pass that `persister` as an option to your Query. This can be done either by passing it to the `defaultOptions` of the `QueryClient` or to any `useQuery` hook instance.
   - If you pass this `persister` as `defaultOptions`, all queries will be persisted to the provided `storage`. You can additionally narrow this down by passing `filters`. In contrast to the `persistClient` plugin, this will not persist the whole query client as a single item, but each query separately. As a key, the query hash is used.
@@ -44,16 +44,18 @@ Garbage collecting a Query from memory **does not** affect the persisted data. T
 
 ```tsx
 import { QueryClient } from '@tanstack/vue-query'
-import { experimental_createPersister } from '@tanstack/query-persist-client-core'
+import { experimental_createQueryPersister } from '@tanstack/query-persist-client-core'
+
+const persister = experimental_createQueryPersister({
+  storage: AsyncStorage,
+  maxAge: 1000 * 60 * 60 * 12, // 12 hours
+})
 
 const queryClient = new QueryClient({
   defaultOptions: {
     queries: {
       gcTime: 1000 * 30, // 30 seconds
-      persister: experimental_createPersister({
-        storage: localStorage,
-        maxAge: 1000 * 60 * 60 * 12, // 12 hours
-      }),
+      persister: persister.persisterFn,
     },
   },
 })
@@ -65,10 +67,10 @@ The `createPersister` plugin technically wraps the `queryFn`, so it doesn't rest
 
 ## API
 
-### `experimental_createPersister`
+### `experimental_createQueryPersister`
 
 ```tsx
-experimental_createPersister(options: StoragePersisterOptions)
+experimental_createQueryPersister(options: StoragePersisterOptions)
 ```
 
 #### `Options`
diff --git a/examples/vue/persister/src/Post.vue b/examples/vue/persister/src/Post.vue
index 43238f97cf..88d8267035 100644
--- a/examples/vue/persister/src/Post.vue
+++ b/examples/vue/persister/src/Post.vue
@@ -4,7 +4,7 @@ import { defineComponent } from 'vue'
 import { useQuery } from '@tanstack/vue-query'
 
 import { Post } from './types'
-import { experimental_createPersister } from '@tanstack/query-persist-client-core'
+import { experimental_createQueryPersister } from '@tanstack/query-persist-client-core'
 
 const fetcher = async (id: number): Promise<Post> =>
   await fetch(`https://jsonplaceholder.typicode.com/posts/${id}`).then(
@@ -24,13 +24,13 @@ export default defineComponent({
     const { isPending, isError, isFetching, data, error } = useQuery({
       queryKey: ['post', props.postId] as const,
       queryFn: () => fetcher(props.postId),
-      persister: experimental_createPersister({
+      persister: experimental_createQueryPersister({
         storage: {
           getItem: (key: string) => get(key),
           setItem: (key: string, value: string) => set(key, value),
           removeItem: (key: string) => del(key),
         },
-      }),
+      }).persisterFn,
     })
 
     return { isPending, isError, isFetching, data, error }
diff --git a/examples/vue/persister/src/Posts.vue b/examples/vue/persister/src/Posts.vue
index 7aa2c02c5b..505a425e27 100644
--- a/examples/vue/persister/src/Posts.vue
+++ b/examples/vue/persister/src/Posts.vue
@@ -1,7 +1,7 @@
 <script lang="ts">
 import { defineComponent } from 'vue'
 import { useQuery } from '@tanstack/vue-query'
-import { experimental_createPersister } from '@tanstack/query-persist-client-core'
+import { experimental_createQueryPersister } from '@tanstack/query-persist-client-core'
 
 import { Post } from './types'
 
@@ -34,9 +34,9 @@ export default defineComponent({
     } = useQuery({
       queryKey: ['posts'] as const,
       queryFn: () => fetcher(),
-      persister: experimental_createPersister({
+      persister: experimental_createQueryPersister({
         storage: localStorage,
-      }),
+      }).persisterFn,
       staleTime: 5000,
     })
 
diff --git a/packages/query-persist-client-core/src/__tests__/createPersister.test.ts b/packages/query-persist-client-core/src/__tests__/createPersister.test.ts
index 9aece2fcd2..d4eea856e4 100644
--- a/packages/query-persist-client-core/src/__tests__/createPersister.test.ts
+++ b/packages/query-persist-client-core/src/__tests__/createPersister.test.ts
@@ -2,7 +2,7 @@ import { describe, expect, test, vi } from 'vitest'
 import { Query, QueryCache, hashKey } from '@tanstack/query-core'
 import {
   PERSISTER_KEY_PREFIX,
-  experimental_createPersister,
+  experimental_createQueryPersister,
 } from '../createPersister'
 import { sleep } from './utils'
 import type { StoragePersisterOptions } from '../createPersister'
@@ -12,10 +12,10 @@ function getFreshStorage() {
   const storage = new Map()
   return {
     getItem: (key: string) => Promise.resolve(storage.get(key)),
-    setItem: async (key: string, value: unknown) => {
+    setItem: (key: string, value: unknown) => {
       storage.set(key, value)
     },
-    removeItem: async (key: string) => {
+    removeItem: (key: string) => {
       storage.delete(key)
     },
   }
@@ -36,7 +36,7 @@ function setupPersister(
 
   const queryFn = vi.fn()
 
-  const persisterFn = experimental_createPersister(persisterOptions)
+  const { persisterFn } = experimental_createQueryPersister(persisterOptions)
 
   const query = new Query({
     cache: new QueryCache(),
@@ -202,7 +202,7 @@ describe('createPersister', () => {
       storageKey,
       JSON.stringify({
         buster: '',
-        state: { dataUpdatedAt },
+        state: { dataUpdatedAt, data: '' },
       }),
     )
 
@@ -231,7 +231,7 @@ describe('createPersister', () => {
       storageKey,
       JSON.stringify({
         buster: '',
-        state: { dataUpdatedAt: Date.now() },
+        state: { dataUpdatedAt: Date.now(), data: '' },
       }),
     )
 
@@ -325,7 +325,7 @@ describe('createPersister', () => {
       storageKey,
       JSON.stringify({
         buster: '',
-        state: { dataUpdatedAt: Date.now() },
+        state: { dataUpdatedAt: Date.now(), data: '' },
       }),
     )
 
diff --git a/packages/query-persist-client-core/src/createPersister.ts b/packages/query-persist-client-core/src/createPersister.ts
index 288e03dc77..a2df7e269a 100644
--- a/packages/query-persist-client-core/src/createPersister.ts
+++ b/packages/query-persist-client-core/src/createPersister.ts
@@ -1,6 +1,7 @@
 import { matchQuery } from '@tanstack/query-core'
 import type {
   Query,
+  QueryClient,
   QueryFilters,
   QueryFunctionContext,
   QueryKey,
@@ -20,6 +21,7 @@ export interface AsyncStorage<TStorageValue = string> {
   getItem: (key: string) => MaybePromise<TStorageValue | undefined | null>
   setItem: (key: string, value: TStorageValue) => MaybePromise<unknown>
   removeItem: (key: string) => MaybePromise<void>
+  entries?: () => MaybePromise<Array<[key: string, value: TStorageValue]>>
 }
 
 export interface StoragePersisterOptions<TStorageValue = string> {
@@ -78,7 +80,7 @@ export const PERSISTER_KEY_PREFIX = 'tanstack-query'
    })
    ```
  */
-export function experimental_createPersister<TStorageValue = string>({
+export function experimental_createQueryPersister<TStorageValue = string>({
   storage,
   buster = '',
   maxAge = 1000 * 60 * 60 * 24,
@@ -91,45 +93,42 @@ export function experimental_createPersister<TStorageValue = string>({
   prefix = PERSISTER_KEY_PREFIX,
   filters,
 }: StoragePersisterOptions<TStorageValue>) {
-  return async function persisterFn<T, TQueryKey extends QueryKey>(
-    queryFn: (context: QueryFunctionContext<TQueryKey>) => T | Promise<T>,
-    context: QueryFunctionContext<TQueryKey>,
-    query: Query,
-  ) {
-    const storageKey = `${prefix}-${query.queryHash}`
-    const matchesFilter = filters ? matchQuery(filters, query) : true
+  function isExpiredOrBusted(persistedQuery: PersistedQuery) {
+    if (persistedQuery.state.dataUpdatedAt) {
+      const queryAge = Date.now() - persistedQuery.state.dataUpdatedAt
+      const expired = queryAge > maxAge
+      const busted = persistedQuery.buster !== buster
 
-    // Try to restore only if we do not have any data in the cache and we have persister defined
-    if (matchesFilter && query.state.data === undefined && storage != null) {
+      if (expired || busted) {
+        return true
+      }
+
+      return false
+    }
+
+    return true
+  }
+
+  async function retrieveQuery<T>(
+    queryHash: string,
+    afterRestoreMacroTask?: (persistedQuery: PersistedQuery) => void,
+  ) {
+    if (storage != null) {
+      const storageKey = `${prefix}-${queryHash}`
       try {
         const storedData = await storage.getItem(storageKey)
         if (storedData) {
           const persistedQuery = await deserialize(storedData)
 
-          if (persistedQuery.state.dataUpdatedAt) {
-            const queryAge = Date.now() - persistedQuery.state.dataUpdatedAt
-            const expired = queryAge > maxAge
-            const busted = persistedQuery.buster !== buster
-            if (expired || busted) {
-              await storage.removeItem(storageKey)
-            } else {
+          if (isExpiredOrBusted(persistedQuery)) {
+            await storage.removeItem(storageKey)
+          } else {
+            if (afterRestoreMacroTask) {
               // Just after restoring we want to get fresh data from the server if it's stale
-              setTimeout(() => {
-                // Set proper updatedAt, since resolving in the first pass overrides those values
-                query.setState({
-                  dataUpdatedAt: persistedQuery.state.dataUpdatedAt,
-                  errorUpdatedAt: persistedQuery.state.errorUpdatedAt,
-                })
-
-                if (query.isStale()) {
-                  query.fetch()
-                }
-              }, 0)
-              // We must resolve the promise here, as otherwise we will have `loading` state in the app until `queryFn` resolves
-              return Promise.resolve(persistedQuery.state.data as T)
+              setTimeout(() => afterRestoreMacroTask(persistedQuery), 0)
             }
-          } else {
-            await storage.removeItem(storageKey)
+            // We must resolve the promise here, as otherwise we will have `loading` state in the app until `queryFn` resolves
+            return persistedQuery.state.data as T
           }
         }
       } catch (err) {
@@ -143,24 +142,117 @@ export function experimental_createPersister<TStorageValue = string>({
       }
     }
 
+    return
+  }
+
+  async function persistQuery(query: Query) {
+    if (storage != null) {
+      const storageKey = `${prefix}-${query.queryHash}`
+      storage.setItem(
+        storageKey,
+        await serialize({
+          state: query.state,
+          queryKey: query.queryKey,
+          queryHash: query.queryHash,
+          buster: buster,
+        }),
+      )
+    }
+  }
+
+  async function persisterFn<T, TQueryKey extends QueryKey>(
+    queryFn: (context: QueryFunctionContext<TQueryKey>) => T | Promise<T>,
+    ctx: QueryFunctionContext<TQueryKey>,
+    query: Query,
+  ) {
+    const matchesFilter = filters ? matchQuery(filters, query) : true
+
+    // Try to restore only if we do not have any data in the cache and we have persister defined
+    if (matchesFilter && query.state.data === undefined && storage != null) {
+      const restoredData = await retrieveQuery(
+        query.queryHash,
+        (persistedQuery: PersistedQuery) => {
+          // Set proper updatedAt, since resolving in the first pass overrides those values
+          query.setState({
+            dataUpdatedAt: persistedQuery.state.dataUpdatedAt,
+            errorUpdatedAt: persistedQuery.state.errorUpdatedAt,
+          })
+
+          if (query.isStale()) {
+            query.fetch()
+          }
+        },
+      )
+
+      if (restoredData != null) {
+        return Promise.resolve(restoredData as T)
+      }
+    }
+
     // If we did not restore, or restoration failed - fetch
-    const queryFnResult = await queryFn(context)
+    const queryFnResult = await queryFn(ctx)
 
     if (matchesFilter && storage != null) {
       // Persist if we have storage defined, we use timeout to get proper state to be persisted
-      setTimeout(async () => {
-        storage.setItem(
-          storageKey,
-          await serialize({
-            state: query.state,
-            queryKey: query.queryKey,
-            queryHash: query.queryHash,
-            buster: buster,
-          }),
-        )
+      setTimeout(() => {
+        persistQuery(query)
       }, 0)
     }
 
     return Promise.resolve(queryFnResult)
   }
+
+  async function persisterGc() {
+    if (storage?.entries) {
+      const entries = await storage.entries()
+      for (const [key, value] of entries) {
+        if (key.startsWith(prefix)) {
+          const persistedQuery = await deserialize(value)
+
+          if (isExpiredOrBusted(persistedQuery)) {
+            await storage.removeItem(key)
+          }
+        }
+      }
+    } else if (process.env.NODE_ENV === 'development') {
+      throw new Error(
+        'Provided storage does not implement `entries` method. Garbage collection is not possible without ability to iterate over storage items.',
+      )
+    }
+  }
+
+  async function persisterRestoreAll(queryClient: QueryClient) {
+    if (storage?.entries) {
+      const entries = await storage.entries()
+      for (const [key, value] of entries) {
+        if (key.startsWith(prefix)) {
+          const persistedQuery = await deserialize(value)
+
+          if (isExpiredOrBusted(persistedQuery)) {
+            await storage.removeItem(key)
+          } else {
+            queryClient.setQueryData(
+              persistedQuery.queryKey,
+              persistedQuery.state.data,
+              {
+                updatedAt: persistedQuery.state.dataUpdatedAt,
+              },
+            )
+          }
+        }
+      }
+    } else if (process.env.NODE_ENV === 'development') {
+      throw new Error(
+        'Provided storage does not implement `entries` method. Restoration of all stored entries is not possible without ability to iterate over storage items.',
+      )
+    }
+  }
+
+  return {
+    persisterFn,
+    persistQuery,
+    retrieveQuery,
+    persisterGc,
+    persisterRestoreAll,
+  }
 }
diff --git a/packages/react-query/src/__tests__/fine-grained-persister.test.tsx b/packages/react-query/src/__tests__/fine-grained-persister.test.tsx
index af376cff87..5ab7c1675e 100644
--- a/packages/react-query/src/__tests__/fine-grained-persister.test.tsx
+++ b/packages/react-query/src/__tests__/fine-grained-persister.test.tsx
@@ -4,7 +4,7 @@ import * as React from 'react'
 import { QueryCache, hashKey } from '@tanstack/query-core'
 import {
   PERSISTER_KEY_PREFIX,
-  experimental_createPersister,
+  experimental_createQueryPersister,
 } from '@tanstack/query-persist-client-core'
 import { useQuery } from '..'
 import { createQueryClient, queryKey, renderWithClient, sleep } from './utils'
@@ -21,10 +21,10 @@ describe('fine grained persister', () => {
     const mapStorage = new Map()
     const storage = {
       getItem: (itemKey: string) => Promise.resolve(mapStorage.get(itemKey)),
-      setItem: async (itemKey: string, value: unknown) => {
+      setItem: (itemKey: string, value: unknown) => {
         mapStorage.set(itemKey, value)
       },
-      removeItem: async (itemKey: string) => {
+      removeItem: (itemKey: string) => {
         mapStorage.delete(itemKey)
       },
     }
@@ -48,9 +48,9 @@ describe('fine grained persister', () => {
       const { data } = useQuery({
         queryKey: key,
         queryFn: spy,
-        persister: experimental_createPersister({
+        persister: experimental_createQueryPersister({
           storage,
-        }),
+        }).persisterFn,
         staleTime: 5000,
       })
 
@@ -75,10 +75,10 @@ describe('fine grained persister', () => {
     const mapStorage = new Map()
     const storage = {
       getItem: (itemKey: string) => Promise.resolve(mapStorage.get(itemKey)),
-      setItem: async (itemKey: string, value: unknown) => {
+      setItem: (itemKey: string, value: unknown) => {
         mapStorage.set(itemKey, value)
       },
-      removeItem: async (itemKey: string) => {
+      removeItem: (itemKey: string) => {
         mapStorage.delete(itemKey)
       },
     }
@@ -102,9 +102,9 @@ describe('fine grained persister', () => {
       const { data } = useQuery({
         queryKey: key,
         queryFn: spy,
-        persister: experimental_createPersister({
+        persister: experimental_createQueryPersister({
           storage,
-        }),
+        }).persisterFn,
       })
 
       return <div ref={(value) => setRef(value)}>{data}</div>
@@ -125,10 +125,10 @@ describe('fine grained persister', () => {
     const mapStorage = new Map()
     const storage = {
       getItem: (itemKey: string) => Promise.resolve(mapStorage.get(itemKey)),
-      setItem: async (itemKey: string, value: unknown) => {
+      setItem: (itemKey: string, value: unknown) => {
         mapStorage.set(itemKey, value)
       },
-      removeItem: async (itemKey: string) => {
+      removeItem: (itemKey: string) => {
         mapStorage.delete(itemKey)
       },
     }
@@ -139,9 +139,9 @@ describe('fine grained persister', () => {
       const { data } = useQuery({
         queryKey: key,
         queryFn: spy,
-        persister: experimental_createPersister({
+        persister: experimental_createQueryPersister({
           storage,
-        }),
+        }).persisterFn,
       })
 
       return <div ref={(value) => setRef(value)}>{data}</div>