Merge pull request #59 from jotaijs/kj/ensure-staletime

bug: add default staletime, setOptions when options change in the optionsAtom
jotaijs · Dec 23, 2023 · fa4aee6 · fa4aee6
2 parents cb1fe8b + 4044c8c
commit fa4aee6
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 61 deletions.
diff --git a/examples/02_typescript_suspense/src/App.tsx b/examples/02_typescript_suspense/src/App.tsx
@@ -1,11 +1,11 @@
-import React from 'react'
+import React, { Suspense } from 'react'
 import { useAtom } from 'jotai/react'
 import { atom } from 'jotai/vanilla'
-import { atomWithQuery } from 'jotai-tanstack-query'
+import { atomWithSuspenseQuery } from 'jotai-tanstack-query'
 
 const idAtom = atom(1)
 
-const userAtom = atomWithQuery<object>((get) => ({
+const userAtom = atomWithSuspenseQuery<object>((get) => ({
   queryKey: ['users', get(idAtom)],
   queryFn: async ({ queryKey: [, id] }) => {
     const res = await fetch(`https://jsonplaceholder.typicode.com/users/${id}`)
@@ -14,10 +14,7 @@ const userAtom = atomWithQuery<object>((get) => ({
 }))
 
 const UserData = () => {
-  const [{ data, isPending, isError }] = useAtom(userAtom)
-
-  if (isPending) return <div>Loading...</div>
-  if (isError) return <div>Error</div>
+  const [{ data }] = useAtom(userAtom)
 
   return <div>{JSON.stringify(data)}</div>
 }
@@ -40,7 +37,9 @@ const Controls = () => {
 const App = () => (
   <>
     <Controls />
-    <UserData />
+    <Suspense fallback="loading">
+      <UserData />
+    </Suspense>
   </>
 )
 

diff --git a/src/atomWithInfiniteQuery.ts b/src/atomWithInfiniteQuery.ts
@@ -24,8 +24,6 @@ export function atomWithInfiniteQuery<
   ) => InfiniteQueryOptions<TQueryFnData, TError, TData, TQueryKey, TPageParam>,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
 ): Atom<InfiniteQueryObserverResult<TData, TError>> {
-  const IN_RENDER = Symbol()
-
   const observerCacheAtom = atom(
     () =>
       new WeakMap<
@@ -46,6 +44,8 @@ export function atomWithInfiniteQuery<
   const optionsAtom = atom((get) => {
     const client = getQueryClient(get)
     const options = getOptions(get)
+    const cache = get(observerCacheAtom)
+    const cachedObserver = cache.get(client)
     const dOptions = client.defaultQueryOptions(
       options
     ) as DefaultedInfiniteQueryObserverOptions<
@@ -59,6 +59,10 @@ export function atomWithInfiniteQuery<
 
     dOptions._optimisticResults = 'optimistic'
 
+    if (cachedObserver) {
+      cachedObserver.setOptions(dOptions, { listeners: false })
+    }
+
     return dOptions
   })
   if (process.env.NODE_ENV !== 'production') {
@@ -70,15 +74,8 @@ export function atomWithInfiniteQuery<
 
     const observerCache = get(observerCacheAtom)
 
-    const observer = observerCache.get(client)
-
-    if (observer) {
-      ;(observer as any)[IN_RENDER] = true
-      observer.setOptions(options, { listeners: false })
-      delete (observer as any)[IN_RENDER]
-
-      return observer
-    }
+    const cachedObserver = observerCache.get(client)
+    if (cachedObserver) return cachedObserver
 
     const newObserver = new InfiniteQueryObserver(client, options)
     observerCache.set(client, newObserver)

diff --git a/src/atomWithQuery.ts b/src/atomWithQuery.ts
@@ -25,8 +25,6 @@ export function atomWithQuery<
   >,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
 ): Atom<QueryObserverResult<TData, TError>> {
-  const IN_RENDER = Symbol()
-
   const observerCacheAtom = atom(
     () =>
       new WeakMap<
@@ -41,10 +39,16 @@ export function atomWithQuery<
   const optionsAtom = atom((get) => {
     const client = getQueryClient(get)
     const options = getOptions(get)
+    const cache = get(observerCacheAtom)
+    const cachedObserver = cache.get(client)
     const dOptions = client.defaultQueryOptions(options)
 
     dOptions._optimisticResults = 'optimistic'
 
+    if (cachedObserver) {
+      cachedObserver.setOptions(dOptions, { listeners: false })
+    }
+
     return dOptions
   })
   if (process.env.NODE_ENV !== 'production') {
@@ -57,15 +61,9 @@ export function atomWithQuery<
 
     const observerCache = get(observerCacheAtom)
 
-    const observer = observerCache.get(client)
+    const cachedObserver = observerCache.get(client)
 
-    if (observer) {
-      ;(observer as any)[IN_RENDER] = true
-      observer.setOptions(options)
-      delete (observer as any)[IN_RENDER]
-
-      return observer
-    }
+    if (cachedObserver) return cachedObserver
 
     const newObserver = new QueryObserver(client, options)
     observerCache.set(client, newObserver)

diff --git a/src/atomWithSuspenseInfiniteQuery.ts b/src/atomWithSuspenseInfiniteQuery.ts
@@ -30,8 +30,6 @@ export const atomWithSuspenseInfiniteQuery = <
   >,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
 ): Atom<Promise<InfiniteQueryObserverSuccessResult<TData, TError>>> => {
-  const IN_RENDER = Symbol()
-
   const observerCacheAtom = atom(
     () =>
       new WeakMap<
@@ -53,6 +51,8 @@ export const atomWithSuspenseInfiniteQuery = <
   const optionsAtom = atom((get) => {
     const client = getQueryClient(get)
     const options = getOptions(get)
+    const cache = get(observerCacheAtom)
+    const cachedObserver = cache.get(client)
     const dOptions = client.defaultQueryOptions(
       options
     ) as DefaultedInfiniteQueryObserverOptions<
@@ -66,6 +66,14 @@ export const atomWithSuspenseInfiniteQuery = <
 
     dOptions._optimisticResults = 'optimistic'
 
+    if (typeof dOptions.staleTime !== 'number') {
+      dOptions.staleTime = 1000
+    }
+
+    if (cachedObserver) {
+      cachedObserver.setOptions(dOptions, { listeners: false })
+    }
+
     return dOptions
   })
   if (process.env.NODE_ENV !== 'production') {
@@ -78,15 +86,9 @@ export const atomWithSuspenseInfiniteQuery = <
 
     const observerCache = get(observerCacheAtom)
 
-    const observer = observerCache.get(client)
-
-    if (observer) {
-      ;(observer as any)[IN_RENDER] = true
-      observer.setOptions(options)
-      delete (observer as any)[IN_RENDER]
+    const cachedObserver = observerCache.get(client)
 
-      return observer
-    }
+    if (cachedObserver) return cachedObserver
 
     const newObserver = new InfiniteQueryObserver(client, options)
     observerCache.set(client, newObserver)

diff --git a/src/atomWithSuspenseQuery.ts b/src/atomWithSuspenseQuery.ts
@@ -28,8 +28,6 @@ export const atomWithSuspenseQuery = <
   | DefinedQueryObserverResult<TData, TError>
   | Promise<DefinedQueryObserverResult<TData, TError>>
 > => {
-  const IN_RENDER = Symbol()
-
   const observerCacheAtom = atom(
     () =>
       new WeakMap<
@@ -44,10 +42,20 @@ export const atomWithSuspenseQuery = <
   const optionsAtom = atom((get) => {
     const client = getQueryClient(get)
     const options = getOptions(get)
+    const cache = get(observerCacheAtom)
+    const cachedObserver = cache.get(client)
     const dOptions = client.defaultQueryOptions(options)
 
     dOptions._optimisticResults = 'optimistic'
 
+    if (typeof dOptions.staleTime !== 'number') {
+      dOptions.staleTime = 1000
+    }
+
+    if (cachedObserver) {
+      cachedObserver.setOptions(dOptions, { listeners: false })
+    }
+
     return dOptions
   })
   if (process.env.NODE_ENV !== 'production') {
@@ -60,15 +68,9 @@ export const atomWithSuspenseQuery = <
 
     const observerCache = get(observerCacheAtom)
 
-    const observer = observerCache.get(client)
-
-    if (observer) {
-      ;(observer as any)[IN_RENDER] = true
-      observer.setOptions(options)
-      delete (observer as any)[IN_RENDER]
+    const cachedObserver = observerCache.get(client)
 
-      return observer
-    }
+    if (cachedObserver) return cachedObserver
 
     const newObserver = new QueryObserver(client, options)
     observerCache.set(client, newObserver)

diff --git a/src/baseAtomWithQuery.ts b/src/baseAtomWithQuery.ts
@@ -139,7 +139,6 @@ export function baseAtomWithQuery<
   ) => QueryObserver<TQueryFnData, TError, TData, TQueryData, TQueryKey>,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
 ) {
-  const IN_RENDER = Symbol()
   const resetAtom = atom(0)
   if (process.env.NODE_ENV !== 'production') {
     resetAtom.debugPrivate = true
@@ -149,13 +148,7 @@ export function baseAtomWithQuery<
     const observer = getObserver(get)
     const source = make<QueryObserverResult<TData, TError>>(({ next }) => {
       const callback = (result: QueryObserverResult<TData, TError>) => {
-        const notifyResult = () => next(result)
-
-        if ((observer as any)[IN_RENDER]) {
-          Promise.resolve().then(notifyResult)
-        } else {
-          notifyResult()
-        }
+        next(result)
       }
 
       return observer.subscribe(callback)
@@ -204,12 +197,11 @@ export function baseAtomWithQuery<
     }
 
     get(resetAtom)
-    const resultAtom = get(dataAtom)
-    const result = get(resultAtom)
+    get(get(dataAtom))
 
-    const optimisticResult = observer.getOptimisticResult(options)
+    const result = observer.getOptimisticResult(options)
 
-    if (shouldSuspend(options, optimisticResult, false)) {
+    if (shouldSuspend(options, result, false)) {
       return observer.fetchOptimistic(options)
     }