Fix refetching behavior for infinite queries

Author: markerikson

packages/toolkit/src/query/core/buildThunks.ts

@@ -25,7 +25,11 @@ import type {
   QueryDefinition,
   ResultTypeFrom,
 } from '../endpointDefinitions'
-import { calculateProvidedBy, isQueryDefinition } from '../endpointDefinitions'
+import {
+  calculateProvidedBy,
+  isInfiniteQueryDefinition,
+  isQueryDefinition,
+} from '../endpointDefinitions'
 import { HandledError } from '../HandledError'
 import type { UnwrapPromise } from '../tsHelpers'
 import type {
@@ -544,11 +548,12 @@ export function buildThunks<
 
         // Start by looking up the existing InfiniteData value from state,
         // falling back to an empty value if it doesn't exist yet
-        const existingData = (getState()[reducerPath].queries[arg.queryCacheKey]
-          ?.data ?? { pages: [], pageParams: [] }) as InfiniteData<
-          unknown,
-          unknown
-        >
+        const blankData = { pages: [], pageParams: [] }
+        const cachedData = getState()[reducerPath].queries[arg.queryCacheKey]
+          ?.data as InfiniteData<unknown, unknown> | undefined
+        const existingData = (
+          isForcedQuery(arg, getState()) || !cachedData ? blankData : cachedData
+        ) as InfiniteData<unknown, unknown>
 
         // If the thunk specified a direction and we do have at least one page,
         // fetch the next or previous page
@@ -564,22 +569,18 @@ export function buildThunks<
           const { initialPageParam = infiniteQueryOptions.initialPageParam } =
             arg as InfiniteQueryThunkArg<any>
 
-          // Fetch first page
-          result = await fetchPage(
-            existingData,
-            existingData.pageParams[0] ?? initialPageParam,
-            maxPages,
-          )
-
-          //original
-          // const remainingPages = pages ?? oldPages.length
-          // const remainingPages = oldPages.length
+          // If we're doing a refetch, we should start from
+          // the first page we have cached.
+          // Otherwise, we should start from the initialPageParam
+          const cachedPageParams = cachedData?.pageParams ?? []
+          const firstPageParam = cachedPageParams[0] ?? initialPageParam
+          const totalPages = cachedPageParams.length
 
-          // TODO This seems pretty wrong
-          const remainingPages = existingData.pages.length
+          // Fetch first page
+          result = await fetchPage(existingData, firstPageParam, maxPages)
 
           // Fetch remaining pages
-          for (let i = 1; i < remainingPages; i++) {
+          for (let i = 1; i < totalPages; i++) {
             const param = getNextPageParam(
               endpointDefinition.infiniteQueryOptions,
               result.data as InfiniteData<unknown, unknown>,
@@ -788,19 +789,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
       }
 
       // if this is forced, continue
-      // if (isForcedQuery(queryThunkArgs, state)) {
-      //   return true
-      // }
-
-      if (
-        isQueryDefinition(endpointDefinition) &&
-        endpointDefinition?.forceRefetch?.({
-          currentArg,
-          previousArg,
-          endpointState: requestState,
-          state,
-        })
-      ) {
+      if (isForcedQuery(queryThunkArgs, state)) {
         return true
       }
 

packages/toolkit/src/query/tests/infiniteQueries.test.ts

@@ -31,18 +31,7 @@ describe('Infinite queries', () => {
     name: string
   }
 
-  server.use(
-    http.get('https://example.com/listItems', ({ request }) => {
-      const url = new URL(request.url)
-      const pageString = url.searchParams.get('page')
-      const pageNum = parseInt(pageString || '0')
-
-      const results: Pokemon[] = [
-        { id: `${pageNum}`, name: `Pokemon ${pageNum}` },
-      ]
-      return HttpResponse.json(results)
-    }),
-  )
+  let counters: Record<string, number> = {}
 
   const pokemonApi = createApi({
     baseQuery: fetchBaseQuery({ baseUrl: 'https://pokeapi.co/api/v2/' }),
@@ -97,9 +86,14 @@ describe('Infinite queries', () => {
           return `https://example.com/listItems?page=${pageParam}`
         },
       }),
-      counter: builder.query<number, string>({
-        queryFn: async () => {
-          return { data: 0 }
+      counters: builder.query<{ id: string; counter: number }, string>({
+        queryFn: async (arg) => {
+          if (!(arg in counters)) {
+            counters[arg] = 0
+          }
+          counters[arg]++
+
+          return { data: { id: arg, counter: counters[arg] } }
         },
       }),
     }),
@@ -114,6 +108,19 @@ describe('Infinite queries', () => {
   )
 
   beforeEach(() => {
+    server.use(
+      http.get('https://example.com/listItems', ({ request }) => {
+        const url = new URL(request.url)
+        const pageString = url.searchParams.get('page')
+        const pageNum = parseInt(pageString || '0')
+
+        const results: Pokemon[] = [
+          { id: `${pageNum}`, name: `Pokemon ${pageNum}` },
+        ]
+        return HttpResponse.json(results)
+      }),
+    )
+
     storeRef = setupApiStore(
       pokemonApi,
       { ...actionsReducer },
@@ -122,6 +129,8 @@ describe('Infinite queries', () => {
       },
     )
 
+    counters = {}
+
     process.env.NODE_ENV = 'development'
   })
 
@@ -218,6 +227,30 @@ describe('Infinite queries', () => {
     }
   })
 
+  test.skip('does not break refetching query endpoints', async () => {
+    const promise0 = storeRef.store.dispatch(
+      pokemonApi.endpoints.counters.initiate('a'),
+    )
+
+    console.log('State after dispatch: ', storeRef.store.getState().api.queries)
+
+    const res0 = await promise0
+
+    console.log('State after promise: ', storeRef.store.getState().api.queries)
+    console.log(storeRef.store.getState().actions)
+
+    const promise1 = storeRef.store.dispatch(
+      pokemonApi.util.upsertQueryData('counters', 'a', { id: 'a', counter: 1 }),
+    )
+
+    console.log('State after dispatch: ', storeRef.store.getState().api.queries)
+
+    const res = await promise1
+
+    console.log('State after promise: ', storeRef.store.getState().api.queries)
+    console.log(storeRef.store.getState().actions)
+  })
+
   test('does not have a page limit without maxPages', async () => {
     for (let i = 1; i <= 10; i++) {
       const res = await storeRef.store.dispatch(
@@ -294,4 +327,80 @@ describe('Infinite queries', () => {
       `getPreviousPageParam for endpoint 'getInfinitePokemon' must be a function if maxPages is used`,
     )
   })
+
+  test('refetches all existing pages', async () => {
+    let hitCounter = 0
+
+    const countersApi = createApi({
+      baseQuery: fakeBaseQuery(),
+      endpoints: (build) => ({
+        counters: build.infiniteQuery<
+          { page: number; hitCounter: number },
+          string,
+          number
+        >({
+          queryFn(page) {
+            hitCounter++
+
+            return { data: { page, hitCounter } }
+          },
+          infiniteQueryOptions: {
+            initialPageParam: 0,
+            getNextPageParam: (
+              lastPage,
+              allPages,
+              lastPageParam,
+              allPageParams,
+            ) => lastPageParam + 1,
+          },
+        }),
+      }),
+    })
+
+    const storeRef = setupApiStore(
+      countersApi,
+      { ...actionsReducer },
+      {
+        withoutTestLifecycles: true,
+      },
+    )
+
+    await storeRef.store.dispatch(
+      countersApi.endpoints.counters.initiate('item', {
+        initialPageParam: 3,
+      }),
+    )
+
+    await storeRef.store.dispatch(
+      countersApi.endpoints.counters.initiate('item', {
+        direction: 'forward',
+      }),
+    )
+
+    const thirdPromise = storeRef.store.dispatch(
+      countersApi.endpoints.counters.initiate('item', {
+        direction: 'forward',
+      }),
+    )
+
+    const thirdRes = await thirdPromise
+    if (thirdRes.status === QueryStatus.fulfilled) {
+      expect(thirdRes.data.pages).toEqual([
+        { page: 3, hitCounter: 1 },
+        { page: 4, hitCounter: 2 },
+        { page: 5, hitCounter: 3 },
+      ])
+    }
+
+    const fourthRes = await thirdPromise.refetch()
+
+    if (fourthRes.status === QueryStatus.fulfilled) {
+      // Refetching should call the query function again for each page
+      expect(fourthRes.data.pages).toEqual([
+        { page: 3, hitCounter: 4 },
+        { page: 4, hitCounter: 5 },
+        { page: 5, hitCounter: 6 },
+      ])
+    }
+  })
 })