feat(core): ensureInfiniteQueryData (#8048)

* feat(core): ensureInfiniteQueryData

* docs: ensureInfiniteQueryData

* feat(types): make sure we can't pass infiniteQueryOptions to non-infinite query functions like fetchQuery

Author: TkDodo

docs/reference/QueryClient.md

@@ -29,6 +29,7 @@ Its available methods are:
 - [`queryClient.prefetchInfiniteQuery`](#queryclientprefetchinfinitequery)
 - [`queryClient.getQueryData`](#queryclientgetquerydata)
 - [`queryClient.ensureQueryData`](#queryclientensurequerydata)
+- [`queryClient.ensureInfiniteQueryData`](#queryclientensureinfinitequerydata)
 - [`queryClient.getQueriesData`](#queryclientgetqueriesdata)
 - [`queryClient.setQueryData`](#queryclientsetquerydata)
 - [`queryClient.getQueryState`](#queryclientgetquerystate)
@@ -200,6 +201,31 @@ const data = await queryClient.ensureQueryData({ queryKey, queryFn })
 
 - `Promise<TData>`
 
+## `queryClient.ensureInfiniteQueryData`
+
+`ensureInfiniteQueryData` is an asynchronous function that can be used to get an existing infinite query's cached data. If the query does not exist, `queryClient.fetchInfiniteQuery` will be called and its results returned.
+
+```tsx
+const data = await queryClient.ensureInfiniteQueryData({
+  queryKey,
+  queryFn,
+  initialPageParam,
+  getNextPageParam,
+})
+```
+
+**Options**
+
+- the same options as [`fetchInfiniteQuery`](#queryclientfetchinfinitequery)
+- `revalidateIfStale: boolean`
+  - Optional
+  - Defaults to `false`
+  - If set to `true`, stale data will be refetched in the background, but cached data will be returned immediately.
+
+**Returns**
+
+- `Promise<InfiniteData<TData, TPageParam>>`
+
 ## `queryClient.getQueriesData`
 
 `getQueriesData` is a synchronous function that can be used to get the cached data of multiple queries. Only queries that match the passed queryKey or queryFilter will be returned. If there are no matching queries, an empty array will be returned.

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

@@ -479,6 +479,69 @@ describe('queryClient', () => {
     })
   })
 
+  describe('ensureInfiniteQueryData', () => {
+    test('should return the cached query data if the query is found', async () => {
+      const key = queryKey()
+      const queryFn = () => Promise.resolve('data')
+
+      queryClient.setQueryData([key, 'id'], { pages: ['bar'], pageParams: [0] })
+
+      await expect(
+        queryClient.ensureInfiniteQueryData({
+          queryKey: [key, 'id'],
+          queryFn,
+          initialPageParam: 1,
+          getNextPageParam: () => undefined,
+        }),
+      ).resolves.toEqual({ pages: ['bar'], pageParams: [0] })
+    })
+
+    test('should fetch the query and return its results if the query is not found', async () => {
+      const key = queryKey()
+      const queryFn = () => Promise.resolve('data')
+
+      await expect(
+        queryClient.ensureInfiniteQueryData({
+          queryKey: [key, 'id'],
+          queryFn,
+          initialPageParam: 1,
+          getNextPageParam: () => undefined,
+        }),
+      ).resolves.toEqual({ pages: ['data'], pageParams: [1] })
+    })
+
+    test('should return the cached query data if the query is found and preFetchQuery in the background when revalidateIfStale is set', async () => {
+      const TIMEOUT = 10
+      const key = queryKey()
+      queryClient.setQueryData([key, 'id'], { pages: ['old'], pageParams: [0] })
+
+      const queryFn = () =>
+        new Promise((resolve) => {
+          setTimeout(() => resolve('new'), TIMEOUT)
+        })
+
+      await expect(
+        queryClient.ensureInfiniteQueryData({
+          queryKey: [key, 'id'],
+          queryFn,
+          initialPageParam: 1,
+          getNextPageParam: () => undefined,
+          revalidateIfStale: true,
+        }),
+      ).resolves.toEqual({ pages: ['old'], pageParams: [0] })
+      await sleep(TIMEOUT + 10)
+      await expect(
+        queryClient.ensureInfiniteQueryData({
+          queryKey: [key, 'id'],
+          queryFn,
+          initialPageParam: 1,
+          getNextPageParam: () => undefined,
+          revalidateIfStale: true,
+        }),
+      ).resolves.toEqual({ pages: ['new'], pageParams: [0] })
+    })
+  })
+
   describe('getQueriesData', () => {
     test('should return the query data for all matched queries', () => {
       const key1 = queryKey()

packages/query-core/src/queryClient.ts

@@ -13,13 +13,13 @@ import { focusManager } from './focusManager'
 import { onlineManager } from './onlineManager'
 import { notifyManager } from './notifyManager'
 import { infiniteQueryBehavior } from './infiniteQueryBehavior'
-import type { QueryState } from './query'
 import type {
   CancelOptions,
   DataTag,
   DefaultError,
   DefaultOptions,
   DefaultedQueryObserverOptions,
+  EnsureInfiniteQueryDataOptions,
   EnsureQueryDataOptions,
   FetchInfiniteQueryOptions,
   FetchQueryOptions,
@@ -40,6 +40,7 @@ import type {
   ResetOptions,
   SetDataOptions,
 } from './types'
+import type { QueryState } from './query'
 import type { MutationFilters, QueryFilters, Updater } from './utils'
 
 // TYPES
@@ -385,7 +386,7 @@ export class QueryClient {
       TData,
       TPageParam
     >(options.pages)
-    return this.fetchQuery(options)
+    return this.fetchQuery(options as any)
   }
 
   prefetchInfiniteQuery<
@@ -406,6 +407,31 @@ export class QueryClient {
     return this.fetchInfiniteQuery(options).then(noop).catch(noop)
   }
 
+  ensureInfiniteQueryData<
+    TQueryFnData,
+    TError = DefaultError,
+    TData = TQueryFnData,
+    TQueryKey extends QueryKey = QueryKey,
+    TPageParam = unknown,
+  >(
+    options: EnsureInfiniteQueryDataOptions<
+      TQueryFnData,
+      TError,
+      TData,
+      TQueryKey,
+      TPageParam
+    >,
+  ): Promise<InfiniteData<TData, TPageParam>> {
+    options.behavior = infiniteQueryBehavior<
+      TQueryFnData,
+      TError,
+      TData,
+      TPageParam
+    >(options.pages)
+
+    return this.ensureQueryData(options as any)
+  }
+
   resumePausedMutations(): Promise<unknown> {
     if (onlineManager.isOnline()) {
       return this.#mutationCache.resumePausedMutations()

packages/query-core/src/types.ts

@@ -448,6 +448,7 @@ export interface FetchQueryOptions<
     QueryOptions<TQueryFnData, TError, TData, TQueryKey, TPageParam>,
     'queryKey'
   > {
+  initialPageParam?: never
   /**
    * The time in milliseconds after data is considered stale.
    * If the data is fresh it will be returned from the cache.
@@ -471,6 +472,22 @@ export interface EnsureQueryDataOptions<
   revalidateIfStale?: boolean
 }
 
+export type EnsureInfiniteQueryDataOptions<
+  TQueryFnData = unknown,
+  TError = DefaultError,
+  TData = TQueryFnData,
+  TQueryKey extends QueryKey = QueryKey,
+  TPageParam = unknown,
+> = FetchInfiniteQueryOptions<
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryKey,
+  TPageParam
+> & {
+  revalidateIfStale?: boolean
+}
+
 type FetchInfiniteQueryPages<TQueryFnData = unknown, TPageParam = unknown> =
   | { pages?: never }
   | {
@@ -484,12 +501,15 @@ export type FetchInfiniteQueryOptions<
   TData = TQueryFnData,
   TQueryKey extends QueryKey = QueryKey,
   TPageParam = unknown,
-> = FetchQueryOptions<
-  TQueryFnData,
-  TError,
-  InfiniteData<TData, TPageParam>,
-  TQueryKey,
-  TPageParam
+> = Omit<
+  FetchQueryOptions<
+    TQueryFnData,
+    TError,
+    InfiniteData<TData, TPageParam>,
+    TQueryKey,
+    TPageParam
+  >,
+  'initialPageParam'
 > &
   InitialPageParam<TPageParam> &
   FetchInfiniteQueryPages<TQueryFnData, TPageParam>

packages/react-query/src/__tests__/infiniteQueryOptions.test-d.tsx

@@ -1,8 +1,9 @@
-import { describe, expectTypeOf, it } from 'vitest'
+import { describe, expectTypeOf, it, test } from 'vitest'
 import { QueryClient, dataTagSymbol } from '@tanstack/query-core'
 import { infiniteQueryOptions } from '../infiniteQueryOptions'
 import { useInfiniteQuery } from '../useInfiniteQuery'
 import { useSuspenseInfiniteQuery } from '../useSuspenseInfiniteQuery'
+import { useQuery } from '../useQuery'
 import type { InfiniteData } from '@tanstack/query-core'
 
 describe('queryOptions', () => {
@@ -133,4 +134,22 @@ describe('queryOptions', () => {
       InfiniteData<string, unknown> | undefined
     >()
   })
+
+  test('should not be allowed to be passed to non-infinite query functions', () => {
+    const queryClient = new QueryClient()
+    const options = infiniteQueryOptions({
+      queryKey: ['key'],
+      queryFn: () => Promise.resolve('string'),
+      getNextPageParam: () => 1,
+      initialPageParam: 1,
+    })
+    // @ts-expect-error cannot pass infinite options to non-infinite query functions
+    useQuery(options)
+    // @ts-expect-error cannot pass infinite options to non-infinite query functions
+    queryClient.ensureQueryData(options)
+    // @ts-expect-error cannot pass infinite options to non-infinite query functions
+    queryClient.fetchQuery(options)
+    // @ts-expect-error cannot pass infinite options to non-infinite query functions
+    queryClient.prefetchQuery(options)
+  })
 })