feat: add atomWithInfiniteQuery (#571)

* feat: add atomWithInfiniteQuery

* chore: size snapshot

.size-snapshot.json

@@ -70,9 +70,9 @@
     }
   },
   "query.js": {
-    "bundled": 2753,
-    "minified": 1235,
-    "gzipped": 608,
+    "bundled": 5309,
+    "minified": 2316,
+    "gzipped": 720,
     "treeshaked": {
       "rollup": {
         "code": 105,

src/query.ts

@@ -1,2 +1,11 @@
-export { queryClientAtom } from './query/queryClientAtom'
+export { queryClientAtom, getQueryClientAtom } from './query/queryClientAtom'
 export { atomWithQuery } from './query/atomWithQuery'
+export { atomWithInfiniteQuery } from './query/atomWithInfiniteQuery'
+export type {
+  AtomWithQueryAction,
+  AtomWithQueryOptions,
+} from './query/atomWithQuery'
+export type {
+  AtomWithInfiniteQueryOptions,
+  AtomWithInfiniteQueryAction,
+} from './query/atomWithInfiniteQuery'

src/query/atomWithInfiniteQuery.ts

@@ -0,0 +1,164 @@
+import {
+  QueryKey,
+  InfiniteQueryObserver,
+  InfiniteQueryObserverOptions,
+  InfiniteData,
+  InitialDataFunction,
+  QueryObserverResult,
+} from 'react-query'
+import { atom } from 'jotai'
+import type { WritableAtom, Getter } from 'jotai'
+import { getQueryClientAtom } from './queryClientAtom'
+
+export type AtomWithInfiniteQueryAction = {
+  type: 'refetch' | 'fetchNextPage' | 'fetchPreviousPage'
+}
+
+export type AtomWithInfiniteQueryOptions<
+  TQueryFnData,
+  TError,
+  TData,
+  TQueryData
+> = InfiniteQueryObserverOptions<TQueryFnData, TError, TData, TQueryData> & {
+  queryKey: QueryKey
+}
+
+export function atomWithInfiniteQuery<
+  TQueryFnData,
+  TError,
+  TData = TQueryFnData,
+  TQueryData = TQueryFnData
+>(
+  createQuery:
+    | AtomWithInfiniteQueryOptions<TQueryFnData, TError, TData, TQueryData>
+    | ((
+        get: Getter
+      ) => AtomWithInfiniteQueryOptions<
+        TQueryFnData,
+        TError,
+        TData,
+        TQueryData
+      >),
+  equalityFn: (
+    a: InfiniteData<TData>,
+    b: InfiniteData<TData>
+  ) => boolean = Object.is
+): WritableAtom<InfiniteData<TData | TQueryData>, AtomWithInfiniteQueryAction> {
+  const queryDataAtom = atom(
+    (get) => {
+      const queryClient = get(getQueryClientAtom)
+      const options =
+        typeof createQuery === 'function' ? createQuery(get) : createQuery
+      let settlePromise:
+        | ((data: InfiniteData<TData> | null, err?: TError) => void)
+        | null = null
+      const getInitialData = () =>
+        typeof options.initialData === 'function'
+          ? (
+              options.initialData as InitialDataFunction<
+                InfiniteData<TQueryData>
+              >
+            )()
+          : options.initialData
+      const dataAtom = atom<
+        | InfiniteData<TData | TQueryData>
+        | Promise<InfiniteData<TData | TQueryData>>
+      >(
+        getInitialData() ||
+          new Promise<InfiniteData<TData>>((resolve, reject) => {
+            settlePromise = (data, err) => {
+              if (err) {
+                reject(err)
+              } else {
+                resolve(data as InfiniteData<TData>)
+              }
+            }
+          })
+      )
+      let setData: (
+        data: InfiniteData<TData> | Promise<InfiniteData<TData>>
+      ) => void = () => {
+        throw new Error('atomWithInfiniteQuery: setting data without mount')
+      }
+      let prevData: InfiniteData<TData> | null = null
+
+      const listener = (
+        result:
+          | QueryObserverResult<InfiniteData<TData>, TError>
+          | { data?: undefined; error: TError }
+      ) => {
+        if (result.error) {
+          if (settlePromise) {
+            settlePromise(null, result.error)
+            settlePromise = null
+          } else {
+            setData(Promise.reject<InfiniteData<TData>>(result.error))
+          }
+          return
+        }
+        if (
+          result.data === undefined ||
+          (prevData !== null && equalityFn(prevData, result.data))
+        ) {
+          return
+        }
+        prevData = result.data
+        if (settlePromise) {
+          settlePromise(result.data)
+          settlePromise = null
+        } else {
+          setData(result.data)
+        }
+      }
+
+      const defaultedOptions = queryClient.defaultQueryObserverOptions(options)
+
+      if (typeof defaultedOptions.staleTime !== 'number') {
+        defaultedOptions.staleTime = 1000
+      }
+
+      const observer = new InfiniteQueryObserver(queryClient, defaultedOptions)
+
+      observer
+        .fetchOptimistic(defaultedOptions)
+        .then(listener)
+        .catch((error) => listener({ error }))
+
+      dataAtom.onMount = (update) => {
+        setData = update
+        const unsubscribe = observer?.subscribe(listener)
+        return unsubscribe
+      }
+      return { dataAtom, observer, options }
+    },
+    (get, set, action: AtomWithInfiniteQueryAction) => {
+      const { observer } = get(queryDataAtom)
+      switch (action.type) {
+        case 'refetch': {
+          void observer.refetch()
+          break
+        }
+        case 'fetchPreviousPage': {
+          void observer.fetchPreviousPage()
+          break
+        }
+        case 'fetchNextPage': {
+          void observer.fetchNextPage()
+          break
+        }
+      }
+    }
+  )
+
+  const queryAtom = atom<
+    InfiniteData<TData | TQueryData>,
+    AtomWithInfiniteQueryAction
+  >(
+    (get) => {
+      const { dataAtom } = get(queryDataAtom)
+      return get(dataAtom)
+    },
+    (_get, set, action) => set(queryDataAtom, action) // delegate action
+  )
+  return queryAtom
+}

src/query/atomWithQuery.ts

@@ -9,9 +9,9 @@ import { atom } from 'jotai'
 import type { WritableAtom, Getter } from 'jotai'
 import { getQueryClientAtom } from './queryClientAtom'
 
-type Action = { type: 'refetch' }
+export type AtomWithQueryAction = { type: 'refetch' }
 
-type AtomQueryOptions<TQueryFnData, TError, TData, TQueryData> =
+export type AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData> =
   QueryObserverOptions<TQueryFnData, TError, TData, TQueryData> & {
     queryKey: QueryKey
   }
@@ -23,12 +23,12 @@ export function atomWithQuery<
   TQueryData = TQueryFnData
 >(
   createQuery:
-    | AtomQueryOptions<TQueryFnData, TError, TData, TQueryData>
+    | AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData>
     | ((
         get: Getter
-      ) => AtomQueryOptions<TQueryFnData, TError, TData, TQueryData>),
+      ) => AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData>),
   equalityFn: (a: TData, b: TData) => boolean = Object.is
-): WritableAtom<TData | TQueryData, Action> {
+): WritableAtom<TData | TQueryData, AtomWithQueryAction> {
   const queryDataAtom = atom(
     (get) => {
       const queryClient = get(getQueryClientAtom)
@@ -100,7 +100,7 @@ export function atomWithQuery<
       }
       return { dataAtom, options }
     },
-    (get, set, action: Action) => {
+    (get, set, action: AtomWithQueryAction) => {
       switch (action.type) {
         case 'refetch': {
           const { dataAtom, options } = get(queryDataAtom)
@@ -113,7 +113,7 @@ export function atomWithQuery<
       }
     }
   )
-  const queryAtom = atom<TData | TQueryData, Action>(
+  const queryAtom = atom<TData | TQueryData, AtomWithQueryAction>(
     (get) => {
       const { dataAtom } = get(queryDataAtom)
       return get(dataAtom)

tests/query/atomWithInfiniteQuery.test.tsx

@@ -0,0 +1,100 @@
+import React, { Suspense } from 'react'
+import { fireEvent, render } from '@testing-library/react'
+import { useAtom } from '../../src/'
+import fakeFetch from './fakeFetch'
+import { getTestProvider } from '../testUtils'
+import { atomWithInfiniteQuery } from '../../src/query/atomWithInfiniteQuery'
+
+const Provider = getTestProvider()
+
+it('infinite query basic test', async () => {
+  const countAtom = atomWithInfiniteQuery<
+    { response: { count: number } },
+    void
+  >(() => ({
+    queryKey: 'count1Infinite',
+    queryFn: async (context) => {
+      const count = context.pageParam ? parseInt(context.pageParam) : 0
+      return fakeFetch({ count })
+    },
+  }))
+
+  const Counter: React.FC = () => {
+    const [data] = useAtom(countAtom)
+    return (
+      <>
+        <div>page count: {data.pages.length}</div>
+      </>
+    )
+  }
+
+  const { findByText } = render(
+    <Provider>
+      <Suspense fallback="loading">
+        <Counter />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('loading')
+  await findByText('page count: 1')
+})
+
+it('infinite query next page test', async () => {
+  const mockFetch = jest.fn(fakeFetch)
+  const countAtom = atomWithInfiniteQuery<
+    { response: { count: number } },
+    void
+  >(() => ({
+    queryKey: 'nextPageAtom',
+    queryFn: (context) => {
+      const count = context.pageParam ? parseInt(context.pageParam) : 0
+      return mockFetch({ count })
+    },
+    getNextPageParam: (lastPage) => {
+      const {
+        response: { count },
+      } = lastPage
+      return (count + 1).toString()
+    },
+    getPreviousPageParam: (lastPage) => {
+      const {
+        response: { count },
+      } = lastPage
+      return (count - 1).toString()
+    },
+  }))
+  const Counter: React.FC = () => {
+    const [data, dispatch] = useAtom(countAtom)
+
+    return (
+      <>
+        <div>page count: {data.pages.length}</div>
+        <button onClick={() => dispatch({ type: 'fetchNextPage' })}>
+          next
+        </button>
+        <button onClick={() => dispatch({ type: 'fetchPreviousPage' })}>
+          prev
+        </button>
+      </>
+    )
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <Suspense fallback="loading">
+        <Counter />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('loading')
+  await findByText('page count: 1')
+  expect(mockFetch).toBeCalledTimes(1)
+  fireEvent.click(getByText('next'))
+  expect(mockFetch).toBeCalledTimes(2)
+  await findByText('page count: 2')
+  fireEvent.click(getByText('prev'))
+  expect(mockFetch).toBeCalledTimes(3)
+  await findByText('page count: 3')
+})