Skip to content

Commit

Permalink
allow to auto-reset error in atomWithSuspenseQuery
Browse files Browse the repository at this point in the history
  • Loading branch information
@grzesiek-ds
grzesiek-ds committed Mar 12, 2024
1 parent fae6e21 commit 350335a
Show file tree
Hide file tree
Showing 2 changed files with 190 additions and 71 deletions.
258 changes: 187 additions & 71 deletions src/baseAtomWithQuery.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@ 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 { BaseAtomWithQueryOptions, queryClientAtom } from 'jotai-tanstack-query'
import { ensureStaleTime, getHasError, shouldSuspend, wait } from './utils'

export function baseAtomWithQuery<
TQueryFnData,
Expand All @@ -31,27 +31,14 @@ export function baseAtomWithQuery<
| QueryObserverResult<TData, TError>
| Promise<QueryObserverResult<TData, TError>>
> {
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<
QueryClient,
QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>
>()
)
if (process.env.NODE_ENV !== 'production') {
observerCacheAtom.debugPrivate = true
}

const defaultedOptionsAtom = atom((get) => {
const client = get(clientAtom)
const options = getOptions(get)
Expand All @@ -63,92 +50,221 @@ 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,
})
}

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<QueryObserverResult<TData, TError> | 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<TData, TError>>
| QueryObserverResult<TData, TError>,
[QueryObserverResult<TData, TError>],
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<TData, TError>
| Promise<QueryObserverResult<TData, TError>>,
[],
Promise<void>
>(
(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
}
3 changes: 3 additions & 0 deletions src/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,3 +80,6 @@ export const ensureStaleTime = (

return defaultedOptions
}

export const wait = (timeout: number): Promise<void> =>
new Promise((resolve) => setTimeout(resolve, timeout))

0 comments on commit 350335a

Please sign in to comment.