123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
import { QueryObserver } from 'react-query'import type {  InitialDataFunction,  QueryKey,  QueryObserverOptions,  QueryObserverResult,} from 'react-query'import { atom } from 'jotai'import type { PrimitiveAtom, WritableAtom } from 'jotai'import { queryClientAtom } from './queryClientAtom'import type { CreateQueryOptions, GetQueryClient } from './types'
export type AtomWithQueryAction = { type: 'refetch' }export type AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData> =  QueryObserverOptions<TQueryFnData, TError, TData, TQueryData> & {    queryKey: QueryKey  }export type AtomWithQueryOptionsWithEnabled<  TQueryFnData,  TError,  TData,  TQueryData> = Omit<  AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData>,  'enabled'> & {  enabled: boolean}
export function atomWithQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithQueryOptionsWithEnabled<TQueryFnData, TError, TData, TQueryData>  >,  getQueryClient?: GetQueryClient): WritableAtom<  TData | TQueryData | undefined,  AtomWithQueryAction,  Promise<void>>
export function atomWithQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData>  >,  getQueryClient?: GetQueryClient): WritableAtom<TData | TQueryData, AtomWithQueryAction, Promise<void>>
export function atomWithQuery<  TQueryFnData,  TError,  TData = TQueryFnData,  TQueryData = TQueryFnData>(  createQuery: CreateQueryOptions<    AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryData>  >,  getQueryClient: GetQueryClient = (get) => get(queryClientAtom)): WritableAtom<  TData | TQueryData | undefined,  AtomWithQueryAction,  Promise<void>> {  const queryDataAtom: WritableAtom<    {      dataAtom: PrimitiveAtom<        TData | TQueryData | Promise<TData | TQueryData> | undefined      >      observer: QueryObserver<TQueryFnData, TError, TData, TQueryData>    },    AtomWithQueryAction,    Promise<void>  > = atom(    (get) => {      const queryClient = getQueryClient(get)      const options =        typeof createQuery === 'function' ? createQuery(get) : createQuery
      let settlePromise:        | ((data: TData | undefined, err?: TError) => void)        | null = null
      const getInitialData = () => {        let data: TQueryData | TData | undefined =          queryClient.getQueryData<TData>(options.queryKey)
        if (data === undefined && options.initialData) {          data =            typeof options.initialData === 'function'              ? (options.initialData as InitialDataFunction<TQueryData>)()              : options.initialData        }        return data      }
      const initialData = getInitialData()
      const dataAtom = atom<        TData | TQueryData | Promise<TData | TQueryData> | undefined      >(        initialData === undefined && options.enabled !== false          ? new Promise<TData>((resolve, reject) => {              settlePromise = (data, err) => {                if (err) {                  reject(err)                } else {                  resolve(data as TData)                }              }            })          : initialData      )      let setData: (data: TData | Promise<TData> | undefined) => void = () => {        throw new Error('atomWithQuery: setting data without mount')      }      const listener = (        result:          | QueryObserverResult<TData, TError>          | { data?: undefined; error: TError }      ) => {        if (result.error) {          if (settlePromise) {            settlePromise(undefined, result.error)            settlePromise = null          } else {            setData(Promise.reject<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 QueryObserver(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 }    },    (get, set, action: AtomWithQueryAction) => {      switch (action.type) {        case 'refetch': {          const { dataAtom, observer } = get(queryDataAtom)          set(dataAtom, new Promise<TData>(() => {})) // infinite pending          const p = Promise.resolve()            .then(() => observer.refetch({ cancelRefetch: true }))            .then(() => {})          return p        }        default:          throw new Error('no action')      }    }  )  const queryAtom = atom<    TData | TQueryData | undefined,    AtomWithQueryAction,    Promise<void>  >(    (get) => {      const { dataAtom } = get(queryDataAtom)      return get(dataAtom)    },    (_get, set, action) => set(queryDataAtom, action) // delegate action  )  return queryAtom}