diff --git a/src/query/atomWithQuery.ts b/src/query/atomWithQuery.ts index 5b12e8d28d..4be8367623 100644 --- a/src/query/atomWithQuery.ts +++ b/src/query/atomWithQuery.ts @@ -1,14 +1,16 @@ -import { QueryObserver } from '@tanstack/query-core' +import { QueryClient, QueryObserver } from '@tanstack/query-core' import type { QueryKey, QueryObserverOptions, QueryObserverResult, } from '@tanstack/query-core' import { atom } from 'jotai' -import type { PrimitiveAtom, WritableAtom } from 'jotai' +import type { WritableAtom } from 'jotai' import { queryClientAtom } from './queryClientAtom' import type { CreateQueryOptions, GetQueryClient } from './types' +type Timeout = ReturnType + type AtomWithQueryAction = { type: 'refetch' } @@ -59,11 +61,7 @@ export function atomWithQuery< > >, getQueryClient?: GetQueryClient -): WritableAtom< - TData | TQueryData | undefined, - AtomWithQueryAction, - void | Promise -> +): WritableAtom export function atomWithQuery< TQueryFnData, @@ -76,7 +74,7 @@ export function atomWithQuery< AtomWithQueryOptions >, getQueryClient?: GetQueryClient -): WritableAtom> +): WritableAtom export function atomWithQuery< TQueryFnData, @@ -90,125 +88,101 @@ export function atomWithQuery< >, getQueryClient: GetQueryClient = (get) => get(queryClientAtom) ): WritableAtom> { - type Options = AtomWithQueryOptions< - TQueryFnData, - TError, - TData, - TQueryData, - TQueryKey - > type Result = QueryObserverResult - const observerAtom = atom((get) => { - const queryClient = getQueryClient(get) - const defaultedOptions = queryClient.defaultQueryOptions< - TQueryFnData, - TError, - TData, - TQueryData, - TQueryKey - >() - const observer = new QueryObserver< - TQueryFnData, - TError, - TData, - TQueryData, - TQueryKey - >(queryClient, defaultedOptions) - return observer - }) + // HACK: having state in atom creator function is not ideal + // because, unlike hooks, it's shared by multiple providers. + const previousDataCache = new WeakMap() - const queryDataAtom: WritableAtom< - { - options: Options - resultAtom: PrimitiveAtom> - unsubIfNotMounted: () => void - }, - AtomWithQueryAction, - void | Promise - > = atom( - (get) => { - const queryClient = getQueryClient(get) - const options = - typeof createQuery === 'function' ? createQuery(get) : createQuery - const defaultedOptions = queryClient.defaultQueryOptions(options) - const observer = get(observerAtom) - observer.destroy() - observer.setOptions(defaultedOptions) - const initialResult = observer.getCurrentResult() + const queryDataAtom = atom((get) => { + const queryClient = getQueryClient(get) + const options = + typeof createQuery === 'function' ? createQuery(get) : createQuery + const observer = new QueryObserver(queryClient, options) + const initialResult = observer.getCurrentResult() + if ( + initialResult.data === undefined && + options.keepPreviousData && + previousDataCache.has(queryClient) + ) { + initialResult.data = previousDataCache.get(queryClient) + } - let resolve: ((result: Result) => void) | null = null - const resultAtom = atom>( - initialResult.data === undefined && options.enabled !== false - ? new Promise((r) => { - resolve = r - }) - : initialResult - ) - let setResult: ((result: Result) => void) | null = null - let unsubscribe: (() => void) | null = null - const unsubIfNotMounted = () => { - if (!setResult) { - unsubscribe?.() - unsubscribe = null - } + let resolve: ((result: Result) => void) | null = null + const makePending = () => + new Promise((r) => { + resolve = r + }) + const resultAtom = atom>( + initialResult.data === undefined && options.enabled !== false + ? makePending() + : initialResult + ) + let setResult: ((result: Result) => void) | null = null + const listener = (result: Result) => { + if (result.isFetching || (!result.isError && result.data === undefined)) { + return } - const listener = (result: Result) => { - if ( - result.isFetching || - (!result.isError && result.data === undefined) - ) { - return - } - if (resolve) { - setTimeout(unsubIfNotMounted, 1000) - resolve(result) - resolve = null - } else if (setResult) { - setResult(result) - } else { - throw new Error('setting result without mount') - } + if (!resolve && !setResult) { + throw new Error('setting result without mount') } - if (options.enabled !== false) { - unsubscribe = observer.subscribe(listener) + if (resolve) { + resolve(result) + resolve = null + } + if (setResult) { + setResult(result) + } + if (result.data !== undefined) { + previousDataCache.set(queryClient, result.data) } - resultAtom.onMount = (update) => { - setResult = update - if (options.enabled !== false && !unsubscribe) { + } + let unsubscribe: (() => void) | null = null + let timer: Timeout | undefined + const startQuery = () => { + if (!setResult && unsubscribe) { + clearTimeout(timer) + unsubscribe() + unsubscribe = null + } + if (options.enabled !== false) { + if (setResult) { + observer.refetch({ cancelRefetch: true }).then(setResult) + } else { unsubscribe = observer.subscribe(listener) - listener(observer.getCurrentResult()) - } - return () => { - setResult = null - unsubscribe?.() } } - return { options, resultAtom, unsubIfNotMounted } - }, - (get, set, action) => { - const observer = get(observerAtom) - const { options, resultAtom, unsubIfNotMounted } = get(queryDataAtom) - if (options.enabled === false) { - return + if (!setResult) { + // not mounted yet + timer = setTimeout(() => { + if (unsubscribe) { + unsubscribe() + unsubscribe = null + } + }, 1000) } - switch (action.type) { - case 'refetch': { - set(resultAtom, new Promise(() => {})) // infinite pending - unsubIfNotMounted() - return observer.refetch({ cancelRefetch: true }).then((result) => { - set(resultAtom, result) - }) + } + startQuery() + resultAtom.onMount = (update) => { + setResult = update + if (unsubscribe) { + clearTimeout(timer) + } else { + startQuery() + } + return () => { + setResult = null + if (unsubscribe) { + unsubscribe() + // FIXME why does this fail? + // unsubscribe = null } } } - ) + return { resultAtom, makePending, startQuery } + }) - const queryAtom = atom< - TData | undefined, - AtomWithQueryAction, - void | Promise - >( + const queryAtom = atom( (get) => { const { resultAtom } = get(queryDataAtom) const result = get(resultAtom) @@ -217,7 +191,15 @@ export function atomWithQuery< } return result.data }, - (_get, set, action) => set(queryDataAtom, action) // delegate action + (get, set, action: AtomWithQueryAction) => { + const { resultAtom, makePending, startQuery } = get(queryDataAtom) + switch (action.type) { + case 'refetch': { + set(resultAtom, makePending()) + startQuery() + } + } + } ) return queryAtom diff --git a/tests/query/atomWithQuery.test.tsx b/tests/query/atomWithQuery.test.tsx index 584e821480..e1c56309bc 100644 --- a/tests/query/atomWithQuery.test.tsx +++ b/tests/query/atomWithQuery.test.tsx @@ -1,5 +1,6 @@ import { Component, StrictMode, Suspense, useContext, useState } from 'react' import type { ReactNode } from 'react' +import { QueryClient } from '@tanstack/query-core' import { fireEvent, render } from '@testing-library/react' import { atom, @@ -722,3 +723,43 @@ describe('error handling', () => { await findByText('count: 3') }) }) + +it('query expected QueryCache test', async () => { + const queryClient = new QueryClient() + const countAtom = atomWithQuery( + () => ({ + queryKey: ['count6'], + queryFn: async () => { + return await fakeFetch({ count: 0 }, false, 100) + }, + }), + () => queryClient + ) + const Counter = () => { + const [ + { + response: { count }, + }, + ] = useAtom(countAtom) + + return ( + <> +
count: {count}
+ + ) + } + + const { findByText } = render( + + + + + + + + ) + + await findByText('loading') + await findByText('count: 0') + expect(queryClient.getQueryCache().getAll().length).toBe(1) +})