fix(core): optimistic results reflect correct state (#7114)

* fix(core): optimistic results reflect correct state

when computing the optimistic result, we need to set / reset all properties that the `fetch` action in the query reducer will also set. This fix extracts the logic into a function and calls it in both places

* chore: comment

Author: TkDodo

packages/query-core/src/query.ts

@@ -552,16 +552,8 @@ export class Query<
         case 'fetch':
           return {
             ...state,
-            fetchFailureCount: 0,
-            fetchFailureReason: null,
+            ...fetchState(state.data, this.options),
             fetchMeta: action.meta ?? null,
-            fetchStatus: canFetch(this.options.networkMode)
-              ? 'fetching'
-              : 'paused',
-            ...(state.data === undefined && {
-              error: null,
-              status: 'pending',
-            }),
           }
         case 'success':
           return {
@@ -620,6 +612,27 @@ export class Query<
   }
 }
 
+export function fetchState<
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryKey extends QueryKey,
+>(
+  data: TData | undefined,
+  options: QueryOptions<TQueryFnData, TError, TData, TQueryKey>,
+) {
+  return {
+    fetchFailureCount: 0,
+    fetchFailureReason: null,
+    fetchStatus: canFetch(options.networkMode) ? 'fetching' : 'paused',
+    ...(data === undefined &&
+      ({
+        error: null,
+        status: 'pending',
+      } as const)),
+  } as const
+}
+
 function getDefaultState<
   TQueryFnData,
   TError,

packages/query-core/src/queryObserver.ts

@@ -9,9 +9,9 @@ import {
 import { notifyManager } from './notifyManager'
 import { focusManager } from './focusManager'
 import { Subscribable } from './subscribable'
-import { canFetch } from './retryer'
-import type { QueryClient } from './queryClient'
+import { fetchState } from './query'
 import type { FetchOptions, Query, QueryState } from './query'
+import type { QueryClient } from './queryClient'
 import type {
   DefaultError,
   DefaultedQueryObserverOptions,
@@ -437,7 +437,7 @@ export class QueryObserver<
       : this.#currentQueryInitialState
 
     const { state } = query
-    let { error, errorUpdatedAt, fetchStatus, status } = state
+    let newState = { ...state }
     let isPlaceholderData = false
     let data: TData | undefined
 
@@ -451,31 +451,31 @@ export class QueryObserver<
         mounted && shouldFetchOptionally(query, prevQuery, options, prevOptions)
 
       if (fetchOnMount || fetchOptionally) {
-        fetchStatus = canFetch(query.options.networkMode)
-          ? 'fetching'
-          : 'paused'
-        if (state.data === undefined) {
-          status = 'pending'
+        newState = {
+          ...newState,
+          ...fetchState(state.data, query.options),
         }
       }
       if (options._optimisticResults === 'isRestoring') {
-        fetchStatus = 'idle'
+        newState.fetchStatus = 'idle'
       }
     }
 
+    let { error, errorUpdatedAt, status } = newState
+
     // Select data if needed
-    if (options.select && state.data !== undefined) {
+    if (options.select && newState.data !== undefined) {
       // Memoize select result
       if (
         prevResult &&
-        state.data === prevResultState?.data &&
+        newState.data === prevResultState?.data &&
         options.select === this.#selectFn
       ) {
         data = this.#selectResult
       } else {
         try {
           this.#selectFn = options.select
-          data = options.select(state.data)
+          data = options.select(newState.data)
           data = replaceData(prevResult?.data, data, options)
           this.#selectResult = data
           this.#selectError = null
@@ -486,7 +486,7 @@ export class QueryObserver<
     }
     // Use query data
     else {
-      data = state.data as unknown as TData
+      data = newState.data as unknown as TData
     }
 
     // Show placeholder data if needed
@@ -541,36 +541,36 @@ export class QueryObserver<
       status = 'error'
     }
 
-    const isFetching = fetchStatus === 'fetching'
+    const isFetching = newState.fetchStatus === 'fetching'
     const isPending = status === 'pending'
     const isError = status === 'error'
 
     const isLoading = isPending && isFetching
-    const hasData = state.data !== undefined
+    const hasData = data !== undefined
 
     const result: QueryObserverBaseResult<TData, TError> = {
       status,
-      fetchStatus,
+      fetchStatus: newState.fetchStatus,
       isPending,
       isSuccess: status === 'success',
       isError,
       isInitialLoading: isLoading,
       isLoading,
       data,
-      dataUpdatedAt: state.dataUpdatedAt,
+      dataUpdatedAt: newState.dataUpdatedAt,
       error,
       errorUpdatedAt,
-      failureCount: state.fetchFailureCount,
-      failureReason: state.fetchFailureReason,
-      errorUpdateCount: state.errorUpdateCount,
-      isFetched: state.dataUpdateCount > 0 || state.errorUpdateCount > 0,
+      failureCount: newState.fetchFailureCount,
+      failureReason: newState.fetchFailureReason,
+      errorUpdateCount: newState.errorUpdateCount,
+      isFetched: newState.dataUpdateCount > 0 || newState.errorUpdateCount > 0,
       isFetchedAfterMount:
-        state.dataUpdateCount > queryInitialState.dataUpdateCount ||
-        state.errorUpdateCount > queryInitialState.errorUpdateCount,
+        newState.dataUpdateCount > queryInitialState.dataUpdateCount ||
+        newState.errorUpdateCount > queryInitialState.errorUpdateCount,
       isFetching,
       isRefetching: isFetching && !isPending,
       isLoadingError: isError && !hasData,
-      isPaused: fetchStatus === 'paused',
+      isPaused: newState.fetchStatus === 'paused',
       isPlaceholderData,
       isRefetchError: isError && hasData,
       isStale: isStale(query, options),

packages/react-query/src/__tests__/useQuery.test.tsx

@@ -6375,4 +6375,104 @@ describe('useQuery', () => {
     await waitFor(() => rendered.getByText('status: success'))
     await waitFor(() => rendered.getByText('data: data'))
   })
+
+  it('should return correct optimistic result when fetching after error', async () => {
+    const key = queryKey()
+    const error = new Error('oh no')
+
+    const results: Array<UseQueryResult<string>> = []
+
+    function Page() {
+      const query = useQuery({
+        queryKey: key,
+        queryFn: async () => {
+          await sleep(10)
+          return Promise.reject(error)
+        },
+        retry: false,
+        notifyOnChangeProps: 'all',
+      })
+
+      results.push(query)
+
+      return (
+        <div>
+          <div>
+            status: {query.status}, {query.fetchStatus}
+          </div>
+          <div>error: {query.error?.message}</div>
+        </div>
+      )
+    }
+
+    function App() {
+      const [enabled, setEnabled] = React.useState(true)
+
+      return (
+        <div>
+          <button onClick={() => setEnabled(!enabled)}>toggle</button>
+          {enabled && <Page />}
+        </div>
+      )
+    }
+
+    const rendered = renderWithClient(queryClient, <App />)
+
+    await waitFor(() => rendered.getByText('status: error, idle'))
+
+    fireEvent.click(rendered.getByRole('button', { name: 'toggle' }))
+    fireEvent.click(rendered.getByRole('button', { name: 'toggle' }))
+
+    await waitFor(() => rendered.getByText('status: error, idle'))
+
+    expect(results).toHaveLength(4)
+
+    // initial fetch
+    expect(results[0]).toMatchObject({
+      status: 'pending',
+      fetchStatus: 'fetching',
+      error: null,
+      errorUpdatedAt: 0,
+      errorUpdateCount: 0,
+      isLoading: true,
+      failureCount: 0,
+      failureReason: null,
+    })
+
+    // error state
+    expect(results[1]).toMatchObject({
+      status: 'error',
+      fetchStatus: 'idle',
+      error,
+      errorUpdateCount: 1,
+      isLoading: false,
+      failureCount: 1,
+      failureReason: error,
+    })
+    expect(results[1]?.errorUpdatedAt).toBeGreaterThan(0)
+
+    // refetch, optimistic state, no errors anymore
+    expect(results[2]).toMatchObject({
+      status: 'pending',
+      fetchStatus: 'fetching',
+      error: null,
+      errorUpdateCount: 1,
+      isLoading: true,
+      failureCount: 0,
+      failureReason: null,
+    })
+    expect(results[2]?.errorUpdatedAt).toBeGreaterThan(0)
+
+    // final state
+    expect(results[3]).toMatchObject({
+      status: 'error',
+      fetchStatus: 'idle',
+      error: error,
+      errorUpdateCount: 2,
+      isLoading: false,
+      failureCount: 1,
+      failureReason: error,
+    })
+    expect(results[3]?.errorUpdatedAt).toBeGreaterThan(0)
+  })
 })