feat(hydration): dehydrate.serialize and hydrate.deserialize (#7615)

* add transformData to restore resolved prefetches

* add failing test

* fix test

* remove transformPromise

* fix

* rev

* rm log

* fix test

* add another test

* serialize/deserialize

* chore: docs + minor refactor

* indentation

* chore: prettier

---------

Co-authored-by: Dominik Dorfmeister <office@dorfmeister.cc>

docs/framework/react/guides/advanced-ssr.md

@@ -432,24 +432,23 @@ export default function Posts() {
 
 > Note that you could also `useQuery` instead of `useSuspenseQuery`, and the Promise would still be picked up correctly. However, NextJs won't suspend in that case and the component will render in the `pending` status, which also opts out of server rendering the content.
 
-If you're using non-JSON data types and serialize the query results on the server, you can specify the `hydrate.transformPromise` option to deserialize the data on the client after the promise is resolved, before the data is put into the cache:
+If you're using non-JSON data types and serialize the query results on the server, you can specify the `dehydrate.serializeData` and `hydrate.deserializeData` options to serialize and deserialize the data on each side of the boundary to ensure the data in the cache is the same format both on the server and the client:
 
 ```tsx
 // app/get-query-client.ts
 import { QueryClient, defaultShouldDehydrateQuery } from '@tanstack/react-query'
-import { deserialize } from './transformer'
+import { deserialize, serialize } from './transformer'
 
 export function makeQueryClient() {
   return new QueryClient({
     defaultOptions: {
+      // ...
       hydrate: {
-        /**
-         * Called when the query is rebuilt from a prefetched
-         * promise, before the query data is put into the cache.
-         */
-        transformPromise: (promise) => promise.then(deserialize),
+        deserializeData: deserialize,
+      },
+      dehydrate: {
+        serializeData: serialize,
       },
-      // ...
     },
   })
 }

docs/framework/react/reference/hydration.md

@@ -37,6 +37,7 @@ const dehydratedState = dehydrate(queryClient, {
       - Return `true` to include this query in dehydration, or `false` otherwise
     - Defaults to only including successful queries
     - If you would like to extend the function while retaining the default behavior, import and execute `defaultShouldDehydrateQuery` as part of the return statement
+  - `serializeData?: (data: any) => any` A function to transform (serialize) data during dehydration.
 
 **Returns**
 
@@ -83,6 +84,7 @@ hydrate(queryClient, dehydratedState, options)
     - Optional
     - `mutations: MutationOptions` The default mutation options to use for the hydrated mutations.
     - `queries: QueryOptions` The default query options to use for the hydrated queries.
+    - `deserializeData?: (data: any) => any` A function to transform (deserialize) data before it is put into the cache.
   - `queryClient?: QueryClient`,
     - Use this to use a custom QueryClient. Otherwise, the one from the nearest context will be used.
 

integrations/react-next-15/app/make-query-client.ts

@@ -22,12 +22,13 @@ export function makeQueryClient() {
          * Called when the query is rebuilt from a prefetched
          * promise, before the query data is put into the cache.
          */
-        transformPromise: (promise) => promise.then(tson.deserialize),
+        deserializeData: tson.deserialize,
       },
       queries: {
         staleTime: 60 * 1000,
       },
       dehydrate: {
+        serializeData: tson.serialize,
         shouldDehydrateQuery: (query) =>
           defaultShouldDehydrateQuery(query) ||
           query.state.status === 'pending',

integrations/react-next-15/app/page.tsx

@@ -13,10 +13,10 @@ export default async function Home() {
     queryKey: ['data'],
     queryFn: async () => {
       await sleep(2000)
-      return tson.serialize({
+      return {
         text: 'data from server',
         date: Temporal.PlainDate.from('2024-01-01'),
-      })
+      }
     },
   })
 

packages/query-core/src/__tests__/hydration.test.tsx

@@ -15,6 +15,11 @@ async function fetchData<TData>(value: TData, ms?: number): Promise<TData> {
   return value
 }
 
+async function fetchDate(value: string, ms?: number): Promise<Date> {
+  await sleep(ms || 0)
+  return new Date(value)
+}
+
 describe('dehydration and rehydration', () => {
   test('should work with serializable values', async () => {
     const queryCache = new QueryCache()
@@ -914,21 +919,22 @@ describe('dehydration and rehydration', () => {
       defaultOptions: {
         dehydrate: {
           shouldDehydrateQuery: () => true,
+          serializeData: (data) => data.toISOString(),
         },
       },
     })
 
     const promise = queryClient.prefetchQuery({
       queryKey: ['transformedStringToDate'],
-      queryFn: () => fetchData('2024-01-01T00:00:00.000Z', 20),
+      queryFn: () => fetchDate('2024-01-01T00:00:00.000Z', 20),
     })
     const dehydrated = dehydrate(queryClient)
     expect(dehydrated.queries[0]?.promise).toBeInstanceOf(Promise)
 
     const hydrationClient = createQueryClient({
       defaultOptions: {
         hydrate: {
-          transformPromise: (p) => p.then((d) => new Date(d)),
+          deserializeData: (data) => new Date(data),
         },
       },
     })
@@ -943,4 +949,81 @@ describe('dehydration and rehydration', () => {
 
     queryClient.clear()
   })
+
+  test('should transform query data if promise is already resolved', async () => {
+    const queryClient = createQueryClient({
+      defaultOptions: {
+        dehydrate: {
+          shouldDehydrateQuery: () => true,
+          serializeData: (data) => data.toISOString(),
+        },
+      },
+    })
+
+    const promise = queryClient.prefetchQuery({
+      queryKey: ['transformedStringToDate'],
+      queryFn: () => fetchDate('2024-01-01T00:00:00.000Z', 0),
+    })
+    await sleep(20)
+    const dehydrated = dehydrate(queryClient)
+
+    const hydrationClient = createQueryClient({
+      defaultOptions: {
+        hydrate: {
+          deserializeData: (data) => new Date(data),
+        },
+      },
+    })
+
+    hydrate(hydrationClient, dehydrated)
+    await promise
+    await waitFor(() =>
+      expect(
+        hydrationClient.getQueryData(['transformedStringToDate']),
+      ).toBeInstanceOf(Date),
+    )
+
+    queryClient.clear()
+  })
+
+  test('should overwrite query in cache if hydrated query is newer (with transformation)', async () => {
+    const hydrationClient = createQueryClient({
+      defaultOptions: {
+        hydrate: {
+          deserializeData: (data) => new Date(data),
+        },
+      },
+    })
+    await hydrationClient.prefetchQuery({
+      queryKey: ['date'],
+      queryFn: () => fetchDate('2024-01-01T00:00:00.000Z', 5),
+    })
+
+    // ---
+
+    const queryClient = createQueryClient({
+      defaultOptions: {
+        dehydrate: {
+          shouldDehydrateQuery: () => true,
+          serializeData: (data) => data.toISOString(),
+        },
+      },
+    })
+    await queryClient.prefetchQuery({
+      queryKey: ['date'],
+      queryFn: () => fetchDate('2024-01-02T00:00:00.000Z', 10),
+    })
+    const dehydrated = dehydrate(queryClient)
+
+    // ---
+
+    hydrate(hydrationClient, dehydrated)
+
+    expect(hydrationClient.getQueryData(['date'])).toStrictEqual(
+      new Date('2024-01-02T00:00:00.000Z'),
+    )
+
+    queryClient.clear()
+    hydrationClient.clear()
+  })
 })

packages/query-core/src/hydration.ts

@@ -13,15 +13,20 @@ import type { Query, QueryState } from './query'
 import type { Mutation, MutationState } from './mutation'
 
 // TYPES
+type TransformerFn = (data: any) => any
+function defaultTransformerFn(data: any): any {
+  return data
+}
 
 export interface DehydrateOptions {
+  serializeData?: TransformerFn
   shouldDehydrateMutation?: (mutation: Mutation) => boolean
   shouldDehydrateQuery?: (query: Query) => boolean
 }
 
 export interface HydrateOptions {
   defaultOptions?: {
-    transformPromise?: (promise: Promise<any>) => Promise<any>
+    deserializeData?: TransformerFn
     queries?: QueryOptions
     mutations?: MutationOptions<unknown, DefaultError, unknown, unknown>
   }
@@ -62,13 +67,21 @@ function dehydrateMutation(mutation: Mutation): DehydratedMutation {
 // consuming the de/rehydrated data, typically with useQuery on the client.
 // Sometimes it might make sense to prefetch data on the server and include
 // in the html-payload, but not consume it on the initial render.
-function dehydrateQuery(query: Query): DehydratedQuery {
+function dehydrateQuery(
+  query: Query,
+  serializeData: TransformerFn,
+): DehydratedQuery {
   return {
-    state: query.state,
+    state: {
+      ...query.state,
+      ...(query.state.data !== undefined && {
+        data: serializeData(query.state.data),
+      }),
+    },
     queryKey: query.queryKey,
     queryHash: query.queryHash,
     ...(query.state.status === 'pending' && {
-      promise: query.promise?.catch((error) => {
+      promise: query.promise?.then(serializeData).catch((error) => {
         if (process.env.NODE_ENV !== 'production') {
           console.error(
             `A query that was dehydrated as pending ended up rejecting. [${query.queryHash}]: ${error}; The error will be redacted in production builds`,
@@ -110,10 +123,17 @@ export function dehydrate(
     client.getDefaultOptions().dehydrate?.shouldDehydrateQuery ??
     defaultShouldDehydrateQuery
 
+  const serializeData =
+    options.serializeData ??
+    client.getDefaultOptions().dehydrate?.serializeData ??
+    defaultTransformerFn
+
   const queries = client
     .getQueryCache()
     .getAll()
-    .flatMap((query) => (filterQuery(query) ? [dehydrateQuery(query)] : []))
+    .flatMap((query) =>
+      filterQuery(query) ? [dehydrateQuery(query, serializeData)] : [],
+    )
 
   return { mutations, queries }
 }
@@ -129,6 +149,10 @@ export function hydrate(
 
   const mutationCache = client.getMutationCache()
   const queryCache = client.getQueryCache()
+  const deserializeData =
+    options?.defaultOptions?.deserializeData ??
+    client.getDefaultOptions().hydrate?.deserializeData ??
+    defaultTransformerFn
 
   // eslint-disable-next-line ts/no-unnecessary-condition
   const mutations = (dehydratedState as DehydratedState).mutations || []
@@ -150,13 +174,19 @@ export function hydrate(
   queries.forEach(({ queryKey, state, queryHash, meta, promise }) => {
     let query = queryCache.get(queryHash)
 
+    const data =
+      state.data === undefined ? state.data : deserializeData(state.data)
+
     // Do not hydrate if an existing query exists with newer data
     if (query) {
       if (query.state.dataUpdatedAt < state.dataUpdatedAt) {
         // omit fetchStatus from dehydrated state
         // so that query stays in its current fetchStatus
-        const { fetchStatus: _ignored, ...dehydratedQueryState } = state
-        query.setState(dehydratedQueryState)
+        const { fetchStatus: _ignored, ...serializedState } = state
+        query.setState({
+          ...serializedState,
+          data,
+        })
       }
     } else {
       // Restore query
@@ -173,19 +203,16 @@ export function hydrate(
         // query being stuck in fetching state upon hydration
         {
           ...state,
+          data,
           fetchStatus: 'idle',
         },
       )
     }
 
     if (promise) {
-      const transformPromise =
-        client.getDefaultOptions().hydrate?.transformPromise
-
       // Note: `Promise.resolve` required cause
       // RSC transformed promises are not thenable
-      const initialPromise =
-        transformPromise?.(Promise.resolve(promise)) ?? promise
+      const initialPromise = Promise.resolve(promise).then(deserializeData)
 
       // this doesn't actually fetch - it just creates a retryer
       // which will re-use the passed `initialPromise`