Hydrate existing queries in a useEffect (#5989)

Ephem · TkDodo · web-flow · commit 3cba6d80e2b2 · 2023-09-15T12:08:52.000+02:00
* feat(react-query): improve hydration timing

BREAKING CHANGE:

* hydrate queries that already exist in the cache in an effect instead of in render
* HydrationBoundary no longer hydrates mutations

* fix(react-query): remove mutations from HydrationBoundary options type

---------

Co-authored-by: Dominik Dorfmeister &lt;office@dorfmeister.cc&gt;
diff --git a/packages/react-query/src/HydrationBoundary.tsx b/packages/react-query/src/HydrationBoundary.tsx
@@ -3,11 +3,17 @@ import * as React from 'react'
 
 import { hydrate } from '@tanstack/query-core'
 import { useQueryClient } from './QueryClientProvider'
-import type { HydrateOptions, QueryClient } from '@tanstack/query-core'
+import type {
+  DehydratedState,
+  HydrateOptions,
+  QueryClient,
+} from '@tanstack/query-core'
 
 export interface HydrationBoundaryProps {
   state?: unknown
-  options?: HydrateOptions
+  options?: Omit<HydrateOptions, 'defaultOptions'> & {
+    defaultOptions?: Omit<HydrateOptions['defaultOptions'], 'mutations'>
+  }
   children?: React.ReactNode
   queryClient?: QueryClient
 }
@@ -19,19 +25,83 @@ export const HydrationBoundary = ({
   queryClient,
 }: HydrationBoundaryProps) => {
   const client = useQueryClient(queryClient)
+  const [hydrationQueue, setHydrationQueue] = React.useState<
+    DehydratedState['queries'] | undefined
+  >()
 
   const optionsRef = React.useRef(options)
   optionsRef.current = options
 
-  // Running hydrate again with the same queries is safe,
-  // it wont overwrite or initialize existing queries,
-  // relying on useMemo here is only a performance optimization.
-  // hydrate can and should be run *during* render here for SSR to work properly
+  // This useMemo is for performance reasons only, everything inside it _must_
+  // be safe to run in every render and code here should be read as "in render".
+  //
+  // This code needs to happen during the render phase, because after initial
+  // SSR, hydration needs to happen _before_ children render. Also, if hydrating
+  // during a transition, we want to hydrate as much as is safe in render so
+  // we can prerender as much as possible.
+  //
+  // For any queries that already exist in the cache, we want to hold back on
+  // hydrating until _after_ the render phase. The reason for this is that during
+  // transitions, we don't want the existing queries and observers to update to
+  // the new data on the current page, only _after_ the transition is committed.
+  // If the transition is aborted, we will have hydrated any _new_ queries, but
+  // we throw away the fresh data for any existing ones to avoid unexpectedly
+  // updating the UI.
   React.useMemo(() => {
     if (state) {
-      hydrate(client, state, optionsRef.current)
+      if (typeof state !== 'object') {
+        return
+      }
+
+      const queryCache = client.getQueryCache()
+      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+      const queries = (state as DehydratedState).queries || []
+
+      const newQueries: DehydratedState['queries'] = []
+      const existingQueries: DehydratedState['queries'] = []
+      for (const dehydratedQuery of queries) {
+        const existingQuery = queryCache.get(dehydratedQuery.queryHash)
+
+        if (!existingQuery) {
+          newQueries.push(dehydratedQuery)
+        } else {
+          const hydrationIsNewer =
+            dehydratedQuery.state.dataUpdatedAt >
+            existingQuery.state.dataUpdatedAt
+          const queryAlreadyQueued = hydrationQueue?.find(
+            (query) => query.queryHash === dehydratedQuery.queryHash,
+          )
+
+          if (
+            hydrationIsNewer &&
+            (!queryAlreadyQueued ||
+              dehydratedQuery.state.dataUpdatedAt >
+                queryAlreadyQueued.state.dataUpdatedAt)
+          ) {
+            existingQueries.push(dehydratedQuery)
+          }
+        }
+      }
+
+      if (newQueries.length > 0) {
+        // It's actually fine to call this with queries/state that already exists
+        // in the cache, or is older. hydrate() is idempotent for queries.
+        hydrate(client, { queries: newQueries }, optionsRef.current)
+      }
+      if (existingQueries.length > 0) {
+        setHydrationQueue((prev) =>
+          prev ? [...prev, ...existingQueries] : existingQueries,
+        )
+      }
+    }
+  }, [client, hydrationQueue, state])
+
+  React.useEffect(() => {
+    if (hydrationQueue) {
+      hydrate(client, { queries: hydrationQueue }, optionsRef.current)
+      setHydrationQueue(undefined)
     }
-  }, [client, state])
+  }, [client, hydrationQueue])
 
   return children as React.ReactElement
 }
diff --git a/packages/react-query/src/__tests__/HydrationBoundary.test.tsx b/packages/react-query/src/__tests__/HydrationBoundary.test.tsx
@@ -137,8 +137,8 @@ describe('React hydration', () => {
         queryFn: () => dataQuery(['should change']),
       })
       await intermediateClient.prefetchQuery({
-        queryKey: ['added string'],
-        queryFn: () => dataQuery(['added string']),
+        queryKey: ['added'],
+        queryFn: () => dataQuery(['added']),
       })
       const dehydrated = dehydrate(intermediateClient)
       intermediateClient.clear()
@@ -147,17 +147,121 @@ describe('React hydration', () => {
         <QueryClientProvider client={queryClient}>
           <HydrationBoundary state={dehydrated}>
             <Page queryKey={['string']} />
-            <Page queryKey={['added string']} />
+            <Page queryKey={['added']} />
           </HydrationBoundary>
         </QueryClientProvider>,
       )
 
-      // Existing query data should be overwritten if older,
-      // so this should have changed
+      // Existing observer should not have updated at this point,
+      // as that would indicate a side effect in the render phase
+      rendered.getByText('string')
+      // New query data should be available immediately
+      rendered.getByText('added')
+
       await sleep(10)
+      // After effects phase has had time to run, the observer should have updated
+      expect(rendered.queryByText('string')).toBeNull()
       rendered.getByText('should change')
-      // New query data should be available immediately
-      rendered.getByText('added string')
+
+      queryClient.clear()
+    })
+
+    // When we hydrate in transitions that are later aborted, it could be
+    // confusing to both developers and users if we suddenly updated existing
+    // state on the screen (why did this update when it was not stale, nothing
+    // remounted, I didn't change tabs etc?).
+    // Any queries that does not exist in the cache yet can still be hydrated
+    // since they don't have any observers on the current page that would update.
+    test('should hydrate new but not existing queries if transition is aborted', async () => {
+      const initialDehydratedState = JSON.parse(stringifiedState)
+      const queryCache = new QueryCache()
+      const queryClient = createQueryClient({ queryCache })
+
+      function Page({ queryKey }: { queryKey: [string] }) {
+        const { data } = useQuery({
+          queryKey,
+          queryFn: () => dataQuery(queryKey),
+        })
+        return (
+          <div>
+            <h1>{data}</h1>
+          </div>
+        )
+      }
+
+      const rendered = render(
+        <QueryClientProvider client={queryClient}>
+          <HydrationBoundary state={initialDehydratedState}>
+            <Page queryKey={['string']} />
+          </HydrationBoundary>
+        </QueryClientProvider>,
+      )
+
+      await rendered.findByText('string')
+
+      const intermediateCache = new QueryCache()
+      const intermediateClient = createQueryClient({
+        queryCache: intermediateCache,
+      })
+      await intermediateClient.prefetchQuery({
+        queryKey: ['string'],
+        queryFn: () => dataQuery(['should not change']),
+      })
+      await intermediateClient.prefetchQuery({
+        queryKey: ['added'],
+        queryFn: () => dataQuery(['added']),
+      })
+      const newDehydratedState = dehydrate(intermediateClient)
+      intermediateClient.clear()
+
+      function Thrower() {
+        throw new Promise(() => {
+          // Never resolve
+        })
+
+        // @ts-ignore
+        return null
+      }
+
+      React.startTransition(() => {
+        rendered.rerender(
+          <React.Suspense fallback="loading">
+            <QueryClientProvider client={queryClient}>
+              <HydrationBoundary state={newDehydratedState}>
+                <Page queryKey={['string']} />
+                <Page queryKey={['added']} />
+                <Thrower />
+              </HydrationBoundary>
+            </QueryClientProvider>
+          </React.Suspense>,
+        )
+
+        rendered.getByText('loading')
+      })
+
+      React.startTransition(() => {
+        rendered.rerender(
+          <QueryClientProvider client={queryClient}>
+            <HydrationBoundary state={initialDehydratedState}>
+              <Page queryKey={['string']} />
+              <Page queryKey={['added']} />
+            </HydrationBoundary>
+          </QueryClientProvider>,
+        )
+
+        // This query existed before the transition so it should stay the same
+        rendered.getByText('string')
+        expect(rendered.queryByText('should not change')).toBeNull()
+        // New query data should be available immediately because it was
+        // hydrated in the previous transition, even though the new dehydrated
+        // state did not contain it
+        rendered.getByText('added')
+      })
+
+      await sleep(10)
+      // It should stay the same even after effects have had a chance to run
+      rendered.getByText('string')
+      expect(rendered.queryByText('should not change')).toBeNull()
 
       queryClient.clear()
     })
