123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180
import type {  Client,  OperationContext,  OperationResult,  RequestPolicy,  TypedDocumentNode,} from '@urql/core'import { pipe, skip, subscribe } from 'wonka'import { atom } from 'jotai'import type { Getter, PrimitiveAtom, WritableAtom } from 'jotai'import { clientAtom } from './clientAtom'
type AtomWithQueryAction = {  type: 'reexecute'  opts?: Partial<OperationContext>}
type OperationResultWithData<Data, Variables> = OperationResult<  Data,  Variables> & {  data: Data}
const isOperationResultWithData = <Data, Variables>(  result: OperationResult<Data, Variables>): result is OperationResultWithData<Data, Variables> =>  'data' in result && !result.error
type QueryArgs<Data, Variables extends object> = {  query: TypedDocumentNode<Data, Variables> | string  variables?: Variables  requestPolicy?: RequestPolicy  context?: Partial<OperationContext>}
type QueryArgsWithPause<Data, Variables extends object> = QueryArgs<  Data,  Variables> & {  pause: boolean}
export function atomWithQuery<Data, Variables extends object>(  createQueryArgs: (get: Getter) => QueryArgs<Data, Variables>,  getClient?: (get: Getter) => Client): WritableAtom<OperationResultWithData<Data, Variables>, AtomWithQueryAction>
export function atomWithQuery<Data, Variables extends object>(  createQueryArgs: (get: Getter) => QueryArgsWithPause<Data, Variables>,  getClient?: (get: Getter) => Client): WritableAtom<  OperationResultWithData<Data, Variables> | null,  AtomWithQueryAction>
export function atomWithQuery<Data, Variables extends object>(  createQueryArgs: (get: Getter) => QueryArgs<Data, Variables>,  getClient: (get: Getter) => Client = (get) => get(clientAtom)) {  type ResultAtom = PrimitiveAtom<    | OperationResultWithData<Data, Variables>    | Promise<OperationResultWithData<Data, Variables>>  >  // TODO we should revisit this for a better solution than refAtom  const refAtom = atom(() => ({} as { resultAtom?: ResultAtom }))  const createResultAtom = (    ref: { resultAtom?: ResultAtom },    client: Client,    args: QueryArgs<Data, Variables>,    opts?: Partial<OperationContext>  ) => {    let resolve:      | ((result: OperationResultWithData<Data, Variables>) => void)      | null = null    const resultAtom: ResultAtom = atom<      | OperationResultWithData<Data, Variables>      | Promise<OperationResultWithData<Data, Variables>>    >(      new Promise<OperationResultWithData<Data, Variables>>((r) => {        resolve = r      })    )    ref.resultAtom = resultAtom // FIXME this is too hacky, may not work in some edge cases    let setResult: (      result: OperationResultWithData<Data, Variables>    ) => void = () => {      throw new Error('setting result without mount')    }    const listener = (result: OperationResult<Data, Variables>) => {      if (resultAtom !== ref.resultAtom) {        // New subscription is working, ignoring old one        return      }      if (!isOperationResultWithData(result)) {        throw new Error('result does not have data')      }      if (resolve) {        resolve(result)        resolve = null      } else {        setResult(result)      }    }    client      .query(args.query, args.variables, {        ...(args.requestPolicy && { requestPolicy: args.requestPolicy }),        ...args.context,        ...opts,      })      .toPromise()      .then(listener)      .catch(() => {        // TODO error handling      })    resultAtom.onMount = (update) => {      setResult = update      const subscription = pipe(        client.query(args.query, args.variables, {          ...(args.requestPolicy && { requestPolicy: args.requestPolicy }),          ...args.context,          ...opts,        }),        skip(1), // handled by toPromise        subscribe(listener)      )      return () => subscription.unsubscribe()    }    return resultAtom  }  const queryResultAtom = atom((get) => {    const args = createQueryArgs(get)    if ((args as { pause?: boolean }).pause) {      return null    }    const client = getClient(get)    const resultAtom = createResultAtom(get(refAtom), client, args)    return { resultAtom, client, args }  })  const overwrittenResultAtom = atom<{    oldResultAtom: ResultAtom    newResultAtom: ResultAtom  } | null>(null)  const queryAtom = atom(    (get) => {      const queryResult = get(queryResultAtom)      if (!queryResult) {        return null      }      let { resultAtom } = queryResult      const overwrittenResult = get(overwrittenResultAtom)      if (overwrittenResult && overwrittenResult.oldResultAtom === resultAtom) {        resultAtom = overwrittenResult.newResultAtom      }      return get(resultAtom)    },    (get, set, action: AtomWithQueryAction) => {      switch (action.type) {        case 'reexecute': {          const queryResult = get(queryResultAtom)          if (!queryResult) {            throw new Error('query is paused')          }          const { resultAtom, client, args } = queryResult          set(overwrittenResultAtom, {            oldResultAtom: resultAtom,            newResultAtom: createResultAtom(              get(refAtom),              client,              args,              action.opts            ),          })        }      }    }  )  return queryAtom}