diff --git a/src/baseAtomWithQuery.ts b/src/baseAtomWithQuery.ts index 112ba9b..a8ea147 100644 --- a/src/baseAtomWithQuery.ts +++ b/src/baseAtomWithQuery.ts @@ -3,11 +3,12 @@ import { QueryKey, QueryObserver, QueryObserverResult, + hashKey, } from '@tanstack/query-core' import { Atom, Getter, atom } from 'jotai' import { queryClientAtom } from './queryClientAtom' import { BaseAtomWithQueryOptions } from './types' -import { ensureStaleTime, getHasError, shouldSuspend } from './utils' +import { ensureStaleTime, getHasError, shouldSuspend, wait } from './utils' export function baseAtomWithQuery< TQueryFnData, @@ -31,16 +32,7 @@ export function baseAtomWithQuery< | QueryObserverResult | Promise> > { - const resetAtom = atom(0) - if (process.env.NODE_ENV !== 'production') { - resetAtom.debugPrivate = true - } - const clientAtom = atom(getQueryClient) - if (process.env.NODE_ENV !== 'production') { - clientAtom.debugPrivate = true - } - const observerCacheAtom = atom( () => new WeakMap< @@ -48,10 +40,6 @@ export function baseAtomWithQuery< QueryObserver >() ) - if (process.env.NODE_ENV !== 'production') { - observerCacheAtom.debugPrivate = true - } - const defaultedOptionsAtom = atom((get) => { const client = get(clientAtom) const options = getOptions(get) @@ -63,6 +51,13 @@ export function baseAtomWithQuery< defaultedOptions._optimisticResults = 'optimistic' if (cachedObserver) { + // This is equivalent to this effect + // https://github.com/TanStack/query/blob/main/packages/react-query/src/useBaseQuery.ts#L94 + // but notice that tanstack/react-query is always returning either observer.getOptimisticResult or observer.fetchOptimistic, + // it never uses value from observer.subscribe (and this implementation does). + // + // Yet, I think that it works as expected, + // because if theres an error, then resultAtom is cleared, therefore after remounting, optimisticResultAtom will be called. cachedObserver.setOptions(defaultedOptions, { listeners: false, }) @@ -70,85 +65,207 @@ export function baseAtomWithQuery< return ensureStaleTime(defaultedOptions) }) - if (process.env.NODE_ENV !== 'production') { - defaultedOptionsAtom.debugPrivate = true - } - const observerAtom = atom((get) => { const client = get(clientAtom) - const defaultedOptions = get(defaultedOptionsAtom) - const observerCache = get(observerCacheAtom) - const cachedObserver = observerCache.get(client) - if (cachedObserver) return cachedObserver - - const newObserver = new Observer(client, defaultedOptions) + if (cachedObserver) { + return cachedObserver + } + // Prevent recalculating this atom by getting defaultedOptions only when creating observer for the first time, + // later observer options will get updated via defaultedOptionsAtom. + const defaultedOptions = get(defaultedOptionsAtom) + const newObserver = new Observer< + TQueryFnData, + TError, + TData, + TQueryData, + TQueryKey + >(client, defaultedOptions) + // Is it good approach to use client as WeakMap key? Would observer get ever garbage collected? + // Why not assign it to scoped variable? + // When whole atom is no longer accessible via reference, then observer should also get garbage collected. observerCache.set(client, newObserver) return newObserver }) - if (process.env.NODE_ENV !== 'production') { - observerAtom.debugPrivate = true - } - - const dataAtom = atom((get) => { + const queryResultGetterAtom = atom((get) => { const observer = get(observerAtom) + // When this atom has value other than null, it means that query result + // has been obtained either by optimisticResultAtom or query subscription defined in onMount below. + // Value is set to null when query is resetting. + const queryResultAtom = atom | null>( + null + ) + // Subscribe to query result updates. Note that this happens only when query did not throw promise or error + // (therefore it has to have valid result) + queryResultAtom.onMount = (set) => { + // tanstack/react-query useBaseQuery is wrapping 'set' with + // https://tanstack.com/query/latest/docs/reference/notifyManager#notifymanagerbatchcalls + // not sure is it relevant here. + const unsubscribe = observer.subscribe(set) + // comment and line from tanstack/react-query useBaseQuery: + // Update result to make sure we did not miss any query updates + // between creating the observer and subscribing to it. + observer.updateResult() - const currentResult = observer.getCurrentResult() - - const resultAtom = atom(currentResult) - if (process.env.NODE_ENV !== 'production') { - resultAtom.debugPrivate = true + return () => { + // Is there any case where it makes sense to reset value on queryAtom unmount? + // I think that it will get garbage collected due to WeakMap usage by jotai. + // Doing this could result in returning outdated state in returnAtom. + // set(null); + unsubscribe() + } } - resultAtom.onMount = (set) => { - const unsubscribe = observer.subscribe((state) => { - set(state) - }) - return unsubscribe + return queryResultAtom + }) + // Simple counter which allows to retrigger optimisticResultAtom calculation, after resetting atom. + const optimisticResultResetAtom = atom(0) + const queryKeyAtom = atom((get) => + hashKey(get(defaultedOptionsAtom).queryKey) + ) + // Reset queryResult each time when queryKey has changed. + const resetQueryResultOnQueryKeyChangeAtom = atom( + (get, { setSelf }) => { + get(queryKeyAtom) + + Promise.resolve().then(setSelf) + }, + (get, set) => { + set(get(queryResultGetterAtom), null) + set(optimisticResultResetAtom, get(optimisticResultResetAtom) + 1) } + ) - return resultAtom - }) - if (process.env.NODE_ENV !== 'production') { - dataAtom.debugPrivate = true - } + // Request optimistic result from observer. Decide whether to suspend or not. + // This atom will recalculate only once per atom reset - see returnAtom. + const optimisticResultAtom = atom< + | Promise> + | QueryObserverResult, + [QueryObserverResult], + void + >( + (get, { setSelf }) => { + const observer = get(observerAtom) + const defaultedOptions = get(defaultedOptionsAtom) + // Recalculate atom when query is resetted. + get(optimisticResultResetAtom) - return atom((get) => { - const observer = get(observerAtom) - const defaultedOptions = get(defaultedOptionsAtom) + const result = observer.getOptimisticResult(defaultedOptions) - const client = getQueryClient(get) + if (!shouldSuspend(defaultedOptions, result, false)) { + // Update queryResultAtom with sync result. + Promise.resolve(result).then(setSelf) - resetAtom.onMount = () => { - return () => { - if (observer.getCurrentResult().isError) { - client.resetQueries({ queryKey: observer.getCurrentQuery().queryKey }) - } + return result } - } - get(resetAtom) - get(get(dataAtom)) + return ( + observer + .fetchOptimistic(defaultedOptions) + .then((succeedResult) => { + setSelf(succeedResult) + + return succeedResult + }) + // useSuspenseQuery is catching fetchOptimistic error, and triggers error boundary. + // Later, when error boundary is resetted, it means that useSuspenseQuery is mounting, and it's resetting the query. + // (therefore for useSuspenseQuery, when error boundary is active then query.state.status === 'error'? Not sure is it like this) + // + // Jotai also has to catch error, because if error would be set as atom value, + // then recovering from error is possible only by outside call + // - see https://github.com/jotaijs/jotai-tanstack-query/issues/32#issue-1684399582 + // This is different than jotai-tanstack-query current implementation. + .catch(() => { + // Since fetchOptimistic failed, error has to be thrown, but let's do it same way as tanstack/react-query useBaseQuery. + // Error handling is done in returnAtom, but first, current query result must be obtained. + // + // observer.currentResult() has outdated value, and observer.updateResult() doesn't help, + // (I'm guessing that's due how observer.getOptimisticResult behaves with suspense = true) + // therefore get optimisticResult but without activating query (_optimisticResults: undefined [I think that's how it works]) + const { _optimisticResults, ...clearedDefaultedOptions } = + defaultedOptions + + const erroredResult = observer.getOptimisticResult( + clearedDefaultedOptions + ) - const result = observer.getOptimisticResult(defaultedOptions) + setSelf(erroredResult) - if (shouldSuspend(defaultedOptions, result, false)) { - return observer.fetchOptimistic(defaultedOptions) + return erroredResult + }) + ) + }, + (get, set, result) => { + // Update queryResultAtom. + set(get(queryResultGetterAtom), result) } + ) + // Better name or just return without variable assignment? + const returnAtom = atom< + | QueryObserverResult + | Promise>, + [], + Promise + >( + (get, { setSelf }) => { + const result = get(get(queryResultGetterAtom)) + // If queryResultAtom has null value it means that value has not yet been obtained from observer. + // Therefore read optimisticResultAtom which activates query and decides whether to suspend or not. + if (result === null) { + // Notice that this atom has valid (up to date) value only just after mounting or resetting queryAtom, + // Later, valid value is held in queryResultAtom. + return get(optimisticResultAtom) + } + // At this point, for atomWithSuspenseQuery it's known that fetchOptimistic promise has been resolved, + // but it could be resolved with error, which is not yet handled. + // Same goes for atomWithQuery (but it could throw error only by using custom throwOnError). + // Therefore verify does query has error or not. - if ( - getHasError({ - result, - query: observer.getCurrentQuery(), - throwOnError: defaultedOptions.throwOnError, - }) - ) { - throw result.error + // Theres potential to retrieve pre 0.8 feature of having async query options getter, + // because we only need options.throwOnError there + // which is non-async function, but this only makes sense for atomWithSuspenseQuery - would have to think about it. + get(resetQueryResultOnQueryKeyChangeAtom) + const options = get(defaultedOptionsAtom) + const query = get(observerAtom).getCurrentQuery() + + if ( + getHasError({ + result, + query, + throwOnError: options.throwOnError, + }) + ) { + // Atom reset has to be scheduled just before throwing error, + // because atom needs to somehow recover from error boundary, + // and if left in error state after unmounting component, thats not possible without outside call. + // If needed that outside call is still possible by using useSetAtom(yourAtomWithQuery), + // but now it's easier to maintain and does not require manually resetting query. + + // While using Promise.resolve().then(setSelf) or wait(0), returnAtom was recalculating before error boundary caught error, + // Maybe @dai-shi could check am I not doing some bad things out there. + wait(1).then(setSelf) + + throw result.error + } + // Non-promise result without error - return it. + return result + }, + async (get, set) => { + // Not sure is it a good call to reset query there. + // tanstack/react-query provides resetErrorBoundary callback, which uses react context under the hood. + // This callback does not reset query, instead it sets a value which is later read when useQuery/useSuspenseQuery + // is triggered (after resetting error boundary). + // https://github.com/TanStack/query/blob/main/packages/react-query/src/errorBoundaryUtils.ts#L29 + // Then, if I understand correctly - query goes back to fetching state, but it still has error. + get(observerAtom).getCurrentQuery().reset() + // Reset to initial state and allow to recalculate optimisticResultAtom + set(get(queryResultGetterAtom), null) + set(optimisticResultResetAtom, get(optimisticResultResetAtom) + 1) } + ) - return result - }) + return returnAtom } diff --git a/src/utils.ts b/src/utils.ts index a42a721..52d67d3 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -80,3 +80,6 @@ export const ensureStaleTime = ( return defaultedOptions } + +export const wait = (timeout: number): Promise => + new Promise((resolve) => setTimeout(resolve, timeout))