fix(hydration): overwrite the existing data in the cache if hydrated data is newer

src/core/query.ts

@@ -65,6 +65,10 @@ const enum ActionType {
   Error,
 }
 
+interface SetDataOptions {
+  updatedAt?: number
+}
+
 interface FailedAction {
   type: ActionType.Failed
 }
@@ -78,6 +82,7 @@ interface SuccessAction<TResult> {
   type: ActionType.Success
   data: TResult | undefined
   canFetchMore?: boolean
+  updatedAt?: number
 }
 
 interface ErrorAction<TError> {
@@ -175,7 +180,10 @@ export class Query<TResult, TError> {
     }
   }
 
-  setData(updater: Updater<TResult | undefined, TResult>): void {
+  setData(
+    updater: Updater<TResult | undefined, TResult>,
+    options?: SetDataOptions
+  ): void {
     const prevData = this.state.data
 
     // Get the new data
@@ -199,6 +207,7 @@ export class Query<TResult, TError> {
       type: ActionType.Success,
       data,
       canFetchMore,
+      updatedAt: options?.updatedAt,
     })
   }
 
@@ -630,7 +639,7 @@ export function queryReducer<TResult, TError>(
         isFetching: false,
         isFetchingMore: false,
         canFetchMore: action.canFetchMore,
-        updatedAt: Date.now(),
+        updatedAt: action.updatedAt ?? Date.now(),
         failureCount: 0,
       }
     case ActionType.Error:

src/hydration/hydration.ts

@@ -3,14 +3,14 @@ import { DEFAULT_CACHE_TIME } from '../core/config'
 import type { Query, QueryCache, QueryKey, QueryConfig } from 'react-query'
 
 export interface DehydratedQueryConfig {
-  queryKey: QueryKey
   cacheTime?: number
-  initialData?: unknown
 }
 
 export interface DehydratedQuery {
-  config: DehydratedQueryConfig
+  queryKey: QueryKey
+  data?: unknown
   updatedAt: number
+  config: DehydratedQueryConfig
 }
 
 export interface DehydratedState {
@@ -28,9 +28,8 @@ function dehydrateQuery<TResult, TError = unknown>(
   query: Query<TResult, TError>
 ): DehydratedQuery {
   const dehydratedQuery: DehydratedQuery = {
-    config: {
-      queryKey: query.queryKey,
-    },
+    config: {},
+    queryKey: query.queryKey,
     updatedAt: query.state.updatedAt,
   }
 
@@ -44,7 +43,7 @@ function dehydrateQuery<TResult, TError = unknown>(
     dehydratedQuery.config.cacheTime = query.cacheTime
   }
   if (query.state.data !== undefined) {
-    dehydratedQuery.config.initialData = query.state.data
+    dehydratedQuery.data = query.state.data
   }
 
   return dehydratedQuery
@@ -82,10 +81,22 @@ export function hydrate<TResult>(
   const queries = (dehydratedState as DehydratedState).queries || []
 
   for (const dehydratedQuery of queries) {
-    const queryKey = dehydratedQuery.config.queryKey
+    const queryKey = dehydratedQuery.queryKey
     const queryConfig = dehydratedQuery.config as QueryConfig<TResult>
-    queryConfig.initialFetched = true
-    const query = queryCache.buildQuery(queryKey, queryConfig)
-    query.state.updatedAt = dehydratedQuery.updatedAt
+
+    let query = queryCache.getQuery<TResult>(queryKey)
+
+    if (query) {
+      if (query.state.updatedAt < dehydratedQuery.updatedAt) {
+        query.setData(dehydratedQuery.data as TResult, {
+          updatedAt: dehydratedQuery.updatedAt,
+        })
+      }
+    } else {
+      query = queryCache.buildQuery<TResult>(queryKey, queryConfig)
+      query.setData(dehydratedQuery.data as TResult, {
+        updatedAt: dehydratedQuery.updatedAt,
+      })
+    }
   }
 }

src/hydration/tests/hydration.test.tsx

@@ -43,24 +43,24 @@ describe('dehydration and rehydration', () => {
 
     const fetchDataAfterHydration = jest.fn()
     await hydrationQueryCache.prefetchQuery('string', fetchDataAfterHydration, {
-      staleTime: 100,
+      staleTime: 1000,
     })
     await hydrationQueryCache.prefetchQuery('number', fetchDataAfterHydration, {
-      staleTime: 100,
+      staleTime: 1000,
     })
     await hydrationQueryCache.prefetchQuery(
       'boolean',
       fetchDataAfterHydration,
-      { staleTime: 100 }
+      { staleTime: 1000 }
     )
     await hydrationQueryCache.prefetchQuery('null', fetchDataAfterHydration, {
-      staleTime: 100,
+      staleTime: 1000,
     })
     await hydrationQueryCache.prefetchQuery('array', fetchDataAfterHydration, {
-      staleTime: 100,
+      staleTime: 1000,
     })
     await hydrationQueryCache.prefetchQuery('nested', fetchDataAfterHydration, {
-      staleTime: 100,
+      staleTime: 1000,
     })
     expect(fetchDataAfterHydration).toHaveBeenCalledTimes(0)
 
@@ -134,7 +134,7 @@ describe('dehydration and rehydration', () => {
     // Exact shape is not important here, just that staleTime and cacheTime
     // (and any future other config) is not included in it
     const dehydratedQuery = dehydrated?.queries.find(
-      query => (query?.config?.queryKey as Array<string>)[0] === 'string'
+      query => (query?.queryKey as Array<string>)[0] === 'string'
     )
     expect(dehydratedQuery).toBeTruthy()
     expect(dehydratedQuery?.config.cacheTime).toBe(undefined)
@@ -179,7 +179,7 @@ describe('dehydration and rehydration', () => {
     // This is testing implementation details that can change and are not
     // part of the public API, but is important for keeping the payload small
     const dehydratedQuery = dehydrated?.queries.find(
-      query => (query?.config?.queryKey as Array<string>)[0] === 'string'
+      query => (query?.queryKey as Array<string>)[0] === 'string'
     )
     expect(dehydratedQuery).toBeUndefined()
 
@@ -196,4 +196,52 @@ describe('dehydration and rehydration', () => {
     queryCache.clear({ notify: false })
     hydrationQueryCache.clear({ notify: false })
   })
+
+  test('should not overwrite query in cache if hydrated query is older', async () => {
+    const queryCache = makeQueryCache()
+    await queryCache.prefetchQuery('string', () => fetchData('string-older', 5))
+    const dehydrated = dehydrate(queryCache)
+    const stringified = JSON.stringify(dehydrated)
+
+    // ---
+
+    const parsed = JSON.parse(stringified)
+    const hydrationQueryCache = makeQueryCache()
+    await hydrationQueryCache.prefetchQuery('string', () =>
+      fetchData('string-newer', 5)
+    )
+
+    hydrate(hydrationQueryCache, parsed)
+    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe(
+      'string-newer'
+    )
+
+    queryCache.clear({ notify: false })
+    hydrationQueryCache.clear({ notify: false })
+  })
+
+  test('should overwrite query in cache if hydrated query is newer', async () => {
+    const hydrationQueryCache = makeQueryCache()
+    await hydrationQueryCache.prefetchQuery('string', () =>
+      fetchData('string-older', 5)
+    )
+
+    // ---
+
+    const queryCache = makeQueryCache()
+    await queryCache.prefetchQuery('string', () => fetchData('string-newer', 5))
+    const dehydrated = dehydrate(queryCache)
+    const stringified = JSON.stringify(dehydrated)
+
+    // ---
+
+    const parsed = JSON.parse(stringified)
+    hydrate(hydrationQueryCache, parsed)
+    expect(hydrationQueryCache.getQuery('string')?.state.data).toBe(
+      'string-newer'
+    )
+
+    queryCache.clear({ notify: false })
+    hydrationQueryCache.clear({ notify: false })
+  })
 })

src/hydration/tests/react.test.tsx

@@ -92,7 +92,7 @@ describe('React hydration', () => {
 
       const intermediateCache = makeQueryCache()
       await intermediateCache.prefetchQuery('string', () =>
-        dataQuery('should not change')
+        dataQuery('should change')
       )
       await intermediateCache.prefetchQuery('added string', dataQuery)
       const dehydrated = dehydrate(intermediateCache)
@@ -107,11 +107,11 @@ describe('React hydration', () => {
         </ReactQueryCacheProvider>
       )
 
-      // Existing query data should not be overwritten,
-      // so this should still be the original data
+      // Existing query data should be overwritten if older,
+      // so this should have changed
       await waitForMs(10)
-      rendered.getByText('string')
-      // But new query data should be available immediately
+      rendered.getByText('should change')
+      // New query data should be available immediately
       rendered.getByText('added string')
 
       clientQueryCache.clear({ notify: false })