From d153545a3d8e7a9be0f47da8f382d41b46ddb186 Mon Sep 17 00:00:00 2001
From: kalijonn <43421621+kalijonn@users.noreply.github.com>
Date: Sun, 22 Sep 2024 19:28:52 +0530
Subject: [PATCH] refresh setter to reset error boundary

---
 __tests__/atomWithSuspenseQuery_spec.tsx | 140 +++++++++++++----------
 src/atomWithInfiniteQuery.ts             |   8 +-
 src/atomWithQuery.ts                     |   8 +-
 src/atomWithSuspenseInfiniteQuery.ts     |  10 +-
 src/atomWithSuspenseQuery.ts             |   6 +-
 src/baseAtomWithQuery.ts                 |  51 +++++----
 6 files changed, 125 insertions(+), 98 deletions(-)

diff --git a/__tests__/atomWithSuspenseQuery_spec.tsx b/__tests__/atomWithSuspenseQuery_spec.tsx
index 21db26b..4df2654 100644
--- a/__tests__/atomWithSuspenseQuery_spec.tsx
+++ b/__tests__/atomWithSuspenseQuery_spec.tsx
@@ -1,7 +1,7 @@
 import React, { StrictMode, Suspense } from 'react'
 import { QueryClient } from '@tanstack/query-core'
 import { fireEvent, render } from '@testing-library/react'
-import { atom, useAtom, useSetAtom } from 'jotai'
+import { Provider, atom, useAtom, useSetAtom } from 'jotai'
 import { ErrorBoundary } from 'react-error-boundary'
 import { atomWithSuspenseQuery } from '../src'
 
@@ -306,68 +306,82 @@ describe('error handling', () => {
     resolve()
     await findByText('errored')
   })
-  // it('can recover from error', async () => {
-  //   let count = -1
-  //   let willThrowError = false
-  //   let resolve = () => {}
-  //   const countAtom = atomWithSuspenseQuery(() => ({
-  //     queryKey: ['error test', 'count2'],
-  //     retry: false,
-  //     queryFn: async () => {
-  //       willThrowError = !willThrowError
-  //       ++count
-  //       await new Promise<void>((r) => (resolve = r))
-  //       if (willThrowError) {
-  //         throw new Error('fetch error')
-  //       }
-  //       return { response: { count } }
-  //     },
-  //   }))
-  //   const Counter = () => {
-  //     const [countData] = useAtom(countAtom)
-  //     return (
-  //       <>
-  //         <div>count: {countData.data?.response.count}</div>
-  //         <button onClick={() => countData.refetch()}>refetch</button>
-  //       </>
-  //     )
-  //   }
-  //   const App = () => {
-  //     return (
-  //       <>
-  //         <ErrorBoundary
-  //           FallbackComponent={({ resetErrorBoundary }) => {
-  //             return (
-  //               <>
-  //                 <h1>errored</h1>
-  //                 <button onClick={resetErrorBoundary}>retry</button>
-  //               </>
-  //             )
-  //           }}>
-  //           <Suspense fallback="loading">
-  //             <Counter />
-  //           </Suspense>
-  //         </ErrorBoundary>
-  //       </>
-  //     )
-  //   }
-  //   const { findByText, getByText } = render(<App />)
-  //   await findByText('loading')
-  //   resolve()
-  //   await findByText('errored')
-  //   fireEvent.click(getByText('retry'))
-  //   await findByText('loading')
-  //   resolve()
-  //   await findByText('count: 1')
-  //   fireEvent.click(getByText('refetch'))
-  //   await findByText('loading')
-  //   resolve()
-  //   await findByText('errored')
-  //   fireEvent.click(getByText('retry'))
-  //   await findByText('loading')
-  //   resolve()
-  //   await findByText('count: 3')
-  // })
+  it('can recover from error', async () => {
+    let count = -1
+    let willThrowError = false
+    let resolve = () => {}
+    const countAtom = atomWithSuspenseQuery(() => ({
+      queryKey: ['error test', 'count2'],
+      retry: false,
+      queryFn: async () => {
+        willThrowError = !willThrowError
+        ++count
+        await new Promise<void>((r) => (resolve = r))
+        if (willThrowError) {
+          throw new Error('fetch error')
+        }
+        return { response: { count } }
+      },
+      throwOnError: true,
+    }))
+    const Counter = () => {
+      const [{ data, refetch, error, isFetching }] = useAtom(countAtom)
+
+      if (error && !isFetching) {
+        throw error
+      }
+
+      return (
+        <>
+          <div>count: {data?.response.count}</div>
+          <button onClick={() => refetch()}>refetch</button>
+        </>
+      )
+    }
+
+    const FallbackComponent: React.FC<{ resetErrorBoundary: () => void }> = ({
+      resetErrorBoundary,
+    }) => {
+      const refresh = useSetAtom(countAtom)
+      return (
+        <>
+          <h1>errored</h1>
+          <button
+            onClick={() => {
+              refresh()
+              resetErrorBoundary()
+            }}>
+            retry
+          </button>
+        </>
+      )
+    }
+    const App = () => {
+      return (
+        <Provider>
+          <ErrorBoundary FallbackComponent={FallbackComponent}>
+            <Suspense fallback="loading">
+              <Counter />
+            </Suspense>
+          </ErrorBoundary>
+        </Provider>
+      )
+    }
+    const { findByText, getByText } = render(<App />)
+    await findByText('loading')
+    resolve()
+    await findByText('errored')
+    fireEvent.click(getByText('retry'))
+    await findByText('loading')
+    resolve()
+    await findByText('count: 1')
+    fireEvent.click(getByText('refetch'))
+    resolve()
+    await findByText('errored')
+    fireEvent.click(getByText('retry'))
+    resolve()
+    await findByText('count: 3')
+  })
 })
 
 it('renews the result when the query changes and a non stale cache is available', async () => {
diff --git a/src/atomWithInfiniteQuery.ts b/src/atomWithInfiniteQuery.ts
index 791f035..bddfa42 100644
--- a/src/atomWithInfiniteQuery.ts
+++ b/src/atomWithInfiniteQuery.ts
@@ -6,7 +6,7 @@ import {
   QueryKey,
   QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, Getter } from 'jotai'
+import { Getter, WritableAtom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
 import {
@@ -34,7 +34,7 @@ export function atomWithInfiniteQuery<
     TPageParam
   >,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<AtomWithInfiniteQueryResult<TData, TError>>
+): WritableAtom<AtomWithInfiniteQueryResult<TData, TError>, [], void>
 export function atomWithInfiniteQuery<
   TQueryFnData,
   TError = DefaultError,
@@ -52,7 +52,7 @@ export function atomWithInfiniteQuery<
     TPageParam
   >,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<DefinedAtomWithInfiniteQueryResult<TData, TError>>
+): WritableAtom<DefinedAtomWithInfiniteQueryResult<TData, TError>, [], void>
 export function atomWithInfiniteQuery<
   TQueryFnData,
   TError = DefaultError,
@@ -71,7 +71,7 @@ export function atomWithInfiniteQuery<
     TPageParam
   >,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<AtomWithInfiniteQueryResult<TData, TError>>
+): WritableAtom<AtomWithInfiniteQueryResult<TData, TError>, [], void>
 export function atomWithInfiniteQuery(
   getOptions: (get: Getter) => AtomWithInfiniteQueryOptions,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
diff --git a/src/atomWithQuery.ts b/src/atomWithQuery.ts
index b06128b..6173214 100644
--- a/src/atomWithQuery.ts
+++ b/src/atomWithQuery.ts
@@ -4,7 +4,7 @@ import {
   QueryKey,
   QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, Getter } from 'jotai'
+import { Getter, WritableAtom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
 import {
@@ -25,7 +25,7 @@ export function atomWithQuery<
     get: Getter
   ) => UndefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<AtomWithQueryResult<TData, TError>>
+): WritableAtom<AtomWithQueryResult<TData, TError>, [], void>
 export function atomWithQuery<
   TQueryFnData = unknown,
   TError = DefaultError,
@@ -36,7 +36,7 @@ export function atomWithQuery<
     get: Getter
   ) => DefinedInitialDataOptions<TQueryFnData, TError, TData, TQueryKey>,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<DefinedAtomWithQueryResult<TData, TError>>
+): WritableAtom<DefinedAtomWithQueryResult<TData, TError>, [], void>
 export function atomWithQuery<
   TQueryFnData = unknown,
   TError = DefaultError,
@@ -47,7 +47,7 @@ export function atomWithQuery<
     get: Getter
   ) => AtomWithQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   getQueryClient?: (get: Getter) => QueryClient
-): Atom<AtomWithQueryResult<TData, TError>>
+): WritableAtom<AtomWithQueryResult<TData, TError>, [], void>
 export function atomWithQuery(
   getOptions: (get: Getter) => AtomWithQueryOptions,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
diff --git a/src/atomWithSuspenseInfiniteQuery.ts b/src/atomWithSuspenseInfiniteQuery.ts
index 399cc67..9192e6f 100644
--- a/src/atomWithSuspenseInfiniteQuery.ts
+++ b/src/atomWithSuspenseInfiniteQuery.ts
@@ -6,7 +6,7 @@ import {
   QueryKey,
   QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter, WritableAtom, atom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
 import {
@@ -33,7 +33,7 @@ export function atomWithSuspenseInfiniteQuery<
     TPageParam
   >,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<AtomWithSuspenseInfiniteQueryResult<TData, TError>> {
+): WritableAtom<AtomWithSuspenseInfiniteQueryResult<TData, TError>, [], void> {
   const suspenseOptionsAtom = atom((get) => {
     const options = getOptions(get)
     return {
@@ -48,5 +48,9 @@ export function atomWithSuspenseInfiniteQuery<
     (get: Getter) => get(suspenseOptionsAtom),
     InfiniteQueryObserver as typeof QueryObserver,
     getQueryClient
-  ) as unknown as Atom<AtomWithSuspenseInfiniteQueryResult<TData, TError>>
+  ) as unknown as WritableAtom<
+    AtomWithSuspenseInfiniteQueryResult<TData, TError>,
+    [],
+    void
+  >
 }
diff --git a/src/atomWithSuspenseQuery.ts b/src/atomWithSuspenseQuery.ts
index 8f6c97b..c3da028 100644
--- a/src/atomWithSuspenseQuery.ts
+++ b/src/atomWithSuspenseQuery.ts
@@ -4,7 +4,7 @@ import {
   QueryKey,
   QueryObserver,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter, WritableAtom, atom } from 'jotai'
 import { baseAtomWithQuery } from './baseAtomWithQuery'
 import { queryClientAtom } from './queryClientAtom'
 import {
@@ -23,7 +23,7 @@ export function atomWithSuspenseQuery<
     get: Getter
   ) => AtomWithSuspenseQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<AtomWithSuspenseQueryResult<TData, TError>> {
+): WritableAtom<AtomWithSuspenseQueryResult<TData, TError>, [], void> {
   const suspenseOptions = atom((get) => {
     const options = getOptions(get)
     return {
@@ -38,5 +38,5 @@ export function atomWithSuspenseQuery<
     (get: Getter) => get(suspenseOptions),
     QueryObserver,
     getQueryClient
-  ) as Atom<AtomWithSuspenseQueryResult<TData, TError>>
+  ) as WritableAtom<AtomWithSuspenseQueryResult<TData, TError>, [], void>
 }
diff --git a/src/baseAtomWithQuery.ts b/src/baseAtomWithQuery.ts
index 2aefea8..118d43f 100644
--- a/src/baseAtomWithQuery.ts
+++ b/src/baseAtomWithQuery.ts
@@ -5,7 +5,7 @@ import {
   QueryObserverResult,
   notifyManager,
 } from '@tanstack/query-core'
-import { Atom, Getter, atom } from 'jotai'
+import { Getter, WritableAtom, atom } from 'jotai'
 import { queryClientAtom } from './queryClientAtom'
 import { BaseAtomWithQueryOptions } from './types'
 import { ensureStaleTime, getHasError, shouldSuspend } from './utils'
@@ -28,10 +28,13 @@ export function baseAtomWithQuery<
   >,
   Observer: typeof QueryObserver,
   getQueryClient: (get: Getter) => QueryClient = (get) => get(queryClientAtom)
-): Atom<
+): WritableAtom<
   | QueryObserverResult<TData, TError>
-  | Promise<QueryObserverResult<TData, TError>>
+  | Promise<QueryObserverResult<TData, TError>>,
+  [],
+  void
 > {
+  const refreshAtom = atom(0)
   const clientAtom = atom(getQueryClient)
   if (process.env.NODE_ENV !== 'production') {
     clientAtom.debugPrivate = true
@@ -115,26 +118,32 @@ export function baseAtomWithQuery<
     dataAtom.debugPrivate = true
   }
 
-  return atom((get) => {
-    const observer = get(observerAtom)
-    const defaultedOptions = get(defaultedOptionsAtom)
+  return atom(
+    (get) => {
+      get(refreshAtom)
+      const observer = get(observerAtom)
+      const defaultedOptions = get(defaultedOptionsAtom)
 
-    const result = get(get(dataAtom))
+      const result = get(get(dataAtom))
 
-    if (shouldSuspend(defaultedOptions, result, false)) {
-      return observer.fetchOptimistic(defaultedOptions)
-    }
+      if (shouldSuspend(defaultedOptions, result, false)) {
+        return observer.fetchOptimistic(defaultedOptions)
+      }
 
-    if (
-      getHasError({
-        result,
-        query: observer.getCurrentQuery(),
-        throwOnError: defaultedOptions.throwOnError,
-      })
-    ) {
-      throw result.error
-    }
+      if (
+        getHasError({
+          result,
+          query: observer.getCurrentQuery(),
+          throwOnError: defaultedOptions.throwOnError,
+        })
+      ) {
+        throw result.error
+      }
 
-    return result
-  })
+      return result
+    },
+    (_get, set) => {
+      set(refreshAtom, (c) => c + 1)
+    }
+  )
 }