TanStack · DamianOsipiuk · Jan 25, 2024
diff --git a/packages/query-core/src/infiniteQueryBehavior.ts b/packages/query-core/src/infiniteQueryBehavior.ts
@@ -116,7 +116,7 @@ export function infiniteQueryBehavior<TQueryFnData, TError, TData, TPageParam>(
       }
       if (context.options.persister) {
         context.fetchFn = () => {
-          return context.options.persister?.(
+          return context.options.persister?.persisterFn(
             fetchFn as any,
             {
               queryKey: context.queryKey,

diff --git a/packages/query-core/src/query.ts b/packages/query-core/src/query.ts
@@ -394,7 +394,7 @@ export class Query<
       }
       this.#abortSignalConsumed = false
       if (this.options.persister) {
-        return this.options.persister(
+        return this.options.persister.persisterFn(
           this.options.queryFn,
           queryFnContext as QueryFunctionContext<TQueryKey>,
           this as unknown as Query,

diff --git a/packages/query-core/src/queryClient.ts b/packages/query-core/src/queryClient.ts
@@ -122,7 +122,7 @@ export class QueryClient {
     return this.#queryCache.find({ queryKey })?.state.data
   }
 
-  ensureQueryData<
+  async ensureQueryData<
     TQueryFnData,
     TError = DefaultError,
     TData = TQueryFnData,
@@ -132,9 +132,27 @@ export class QueryClient {
   ): Promise<TData> {
     const cachedData = this.getQueryData<TData>(options.queryKey)
 
-    return cachedData !== undefined
-      ? Promise.resolve(cachedData)
-      : this.fetchQuery(options)
+    if (cachedData !== undefined) {
+      return Promise.resolve(cachedData)
+    }
+
+    // Check persister if defined
+    const defaultedOptions = this.defaultQueryOptions({
+      queryKey: options.queryKey,
+    })
+
+    if (defaultedOptions.persister) {
+      const persistedData =
+        await defaultedOptions.persister.restoreQuery<TData>(
+          options.queryHash || hashKey(options.queryKey),
+        )
+
+      if (persistedData != null) {
+        return Promise.resolve(persistedData)
+      }
+    }
+
+    return this.fetchQuery(options)
   }
 
   getQueriesData<TQueryFnData = unknown>(
@@ -181,9 +199,14 @@ export class QueryClient {
       QueryKey
     >({ queryKey })
 
-    return this.#queryCache
-      .build(this, defaultedOptions)
-      .setData(data, { ...options, manual: true })
+    const newQuery = this.#queryCache.build(this, defaultedOptions)
+    const result = newQuery.setData(data, { ...options, manual: true })
+
+    if (defaultedOptions.persister) {
+      defaultedOptions.persister.persistQuery(newQuery)
+    }
+
+    return result
   }
 
   setQueriesData<TQueryFnData>(

diff --git a/packages/query-core/src/types.ts b/packages/query-core/src/types.ts
@@ -1,7 +1,7 @@
 /* istanbul ignore file */
 
 import type { MutationState } from './mutation'
-import type { FetchDirection, Query, QueryBehavior } from './query'
+import type { FetchDirection, Query, QueryBehavior, QueryState } from './query'
 import type { RetryDelayValue, RetryValue } from './retryer'
 import type { QueryFilters, QueryTypeFilter } from './utils'
 import type { QueryCache } from './queryCache'
@@ -148,11 +148,23 @@ export interface QueryOptions<
    */
   gcTime?: number
   queryFn?: QueryFunction<TQueryFnData, TQueryKey, TPageParam>
-  persister?: QueryPersister<
-    NoInfer<TQueryFnData>,
-    NoInfer<TQueryKey>,
-    NoInfer<TPageParam>
-  >
+  persister?: {
+    persisterFn: QueryPersister<
+      NoInfer<TQueryFnData>,
+      NoInfer<TQueryKey>,
+      NoInfer<TPageParam>
+    >
+    persistQuery: (query: Query) => Promise<void>
+    restoreQuery: <T>(
+      queryHash: string,
+      afterRestoreMacroTask?: (persistedQuery: {
+        buster: string
+        queryHash: string
+        queryKey: QueryKey
+        state: QueryState
+      }) => void,
+    ) => Promise<T | undefined>
+  }
   queryHash?: string
   queryKey?: TQueryKey
   queryKeyHashFn?: QueryKeyHashFunction<TQueryKey>

diff --git a/packages/query-persist-client-core/src/__tests__/createPersister.test.ts b/packages/query-persist-client-core/src/__tests__/createPersister.test.ts
@@ -36,7 +36,7 @@ function setupPersister(
 
   const queryFn = vi.fn()
 
-  const persisterFn = experimental_createPersister(persisterOptions)
+  const { persisterFn } = experimental_createPersister(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
@@ -91,16 +91,12 @@ 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,
+  async function restoreQuery<T>(
+    queryHash: string,
+    afterRestoreMacroTask?: (persistedQuery: PersistedQuery) => void,
   ) {
-    const storageKey = `${prefix}-${query.queryHash}`
-    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) {
+    if (storage != null) {
+      const storageKey = `${prefix}-${queryHash}`
       try {
         const storedData = await storage.getItem(storageKey)
         if (storedData) {
@@ -113,20 +109,12 @@ export function experimental_createPersister<TStorageValue = string>({
             if (expired || busted) {
               await storage.removeItem(storageKey)
             } else {
-              // 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)
+              if (afterRestoreMacroTask) {
+                // Just after restoring we want to get fresh data from the server if it's stale
+                setTimeout(() => afterRestoreMacroTask(persistedQuery), 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)
+              return persistedQuery.state.data as T
             }
           } else {
             await storage.removeItem(storageKey)
@@ -143,24 +131,69 @@ 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 restoreQuery(
+        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)
   }
+
+  return {
+    persisterFn,
+    persistQuery,
+    restoreQuery,
+  }
 }