123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
import { InfiniteQueryObserver, isCancelledError } from 'react-query'import type {  InfiniteData,  InfiniteQueryObserverOptions,  InitialDataFunction,  QueryKey,  QueryObserverResult,  RefetchOptions,  RefetchQueryFilters,} from 'react-query'import { atom } from 'jotai'import type { WritableAtom } from 'jotai'import { queryClientAtom } from './queryClientAtom'import { CreateQueryOptions, GetQueryClient } from './types'
export type AtomWithInfiniteQueryAction<TQueryFnData> =  | ({ type: 'refetch' } & Partial<      RefetchOptions & RefetchQueryFilters<TQueryFnData>    >)  | { type: 'fetchNextPage' }  | { type: 'fetchPreviousPage' }
export type AtomWithInfiniteQueryOptions<  TQueryFnData,  TError,  TData,  TQueryData> = InfiniteQueryObserverOptions<TQueryFnData, TError, TData, TQueryData> & {  queryKey: QueryKey}
export type AtomWithInfiniteQueryOptionsWithEnabled<  TQueryFnData,  TError,  TData,  TQueryData> = Omit<  AtomWithInfiniteQueryOptions<TQueryFnData, TError, TData, TQueryData>,  'enabled'> & {  enabled: boolean}
export function atomWithInfiniteQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithInfiniteQueryOptionsWithEnabled<      TQueryFnData,      TError,      TData,      TQueryData    >  >,  getQueryClient?: GetQueryClient): WritableAtom<  InfiniteData<TData | TQueryData> | undefined,  AtomWithInfiniteQueryAction<TQueryFnData>>
export function atomWithInfiniteQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithInfiniteQueryOptions<TQueryFnData, TError, TData, TQueryData>  >,  getQueryClient?: GetQueryClient): WritableAtom<  InfiniteData<TData | TQueryData>,  AtomWithInfiniteQueryAction<TQueryFnData>>
export function atomWithInfiniteQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithInfiniteQueryOptions<TQueryFnData, TError, TData, TQueryData>  >,  getQueryClient: GetQueryClient = (get) => get(queryClientAtom)): WritableAtom<  InfiniteData<TData | TQueryData> | undefined,  AtomWithInfiniteQueryAction<TQueryFnData>> {  const queryDataAtom = atom(    (get) => {      const queryClient = getQueryClient(get)      const options =        typeof createQuery === 'function' ? createQuery(get) : createQuery      let settlePromise:        | ((data: InfiniteData<TData> | undefined, err?: TError) => void)        | null = null
      const getInitialData = () => {        let data: InfiniteData<TQueryData | TData> | undefined =          queryClient.getQueryData<InfiniteData<TData>>(options.queryKey)
        if (data === undefined && options.initialData) {          data =            typeof options.initialData === 'function'              ? (                  options.initialData as InitialDataFunction<                    InfiniteData<TQueryData>                  >                )()              : options.initialData        }
        return data      }
      const initialData = getInitialData()
      const dataAtom = atom<        | InfiniteData<TData | TQueryData>        | Promise<InfiniteData<TData | TQueryData>>        | undefined      >(        initialData === undefined && options.enabled !== false          ? new Promise<InfiniteData<TData>>((resolve, reject) => {              settlePromise = (data, err) => {                if (err) {                  reject(err)                } else {                  resolve(data as InfiniteData<TData>)                }              }            })          : initialData      )      let setData: (        data: InfiniteData<TData> | Promise<InfiniteData<TData>> | undefined      ) => void = () => {        throw new Error('atomWithInfiniteQuery: setting data without mount')      }      const listener = (        result:          | QueryObserverResult<InfiniteData<TData>, TError>          | { data?: undefined; error: TError }      ) => {        if (result.error && !isCancelledError(result.error)) {          if (settlePromise) {            settlePromise(undefined, result.error)            settlePromise = null          } else {            setData(Promise.reject<InfiniteData<TData>>(result.error))          }          return        }        if (result.data === undefined) {          return        }        if (settlePromise) {          settlePromise(result.data)          settlePromise = null        } else {          setData(result.data)        }      }
      const defaultedOptions = queryClient.defaultQueryObserverOptions(options)      if (initialData === undefined && options.enabled !== false) {        if (typeof defaultedOptions.staleTime !== 'number') {          defaultedOptions.staleTime = 1000        }      }      const observer = new InfiniteQueryObserver(queryClient, defaultedOptions)      if (initialData === undefined && options.enabled !== false) {        observer          .fetchOptimistic(defaultedOptions)          .then(listener)          .catch((error) => listener({ error }))      }
      dataAtom.onMount = (update) => {        setData = update        if (options.enabled !== false) {          return observer.subscribe(listener)        }      }      return { dataAtom, observer, options }    },    (get, _set, action: AtomWithInfiniteQueryAction<TQueryFnData>) => {      const { observer } = get(queryDataAtom)      switch (action.type) {        case 'refetch': {          const { type: _type, ...options } = action          void observer.refetch(options)          break        }        case 'fetchPreviousPage': {          void observer.fetchPreviousPage()          break        }        case 'fetchNextPage': {          void observer.fetchNextPage()          break        }      }    }  )
  const queryAtom = atom<    InfiniteData<TData | TQueryData> | undefined,    AtomWithInfiniteQueryAction<TQueryFnData>  >(    (get) => {      const { dataAtom } = get(queryDataAtom)      return get(dataAtom)    },    (_get, set, action) => set(queryDataAtom, action) // delegate action  )  return queryAtom}