From cbec5d7cead2b639db2cf883b888f6927e795d86 Mon Sep 17 00:00:00 2001
From: Niek <just.niek@gmail.com>
Date: Sat, 6 Feb 2021 10:13:27 +0100
Subject: [PATCH] feat: add suspense to useQueries

---
 src/core/queriesObserver.ts         |  20 +++-
 src/react/tests/useQueries.test.tsx | 167 +++++++++++++++++++++++++++-
 src/react/useQueries.ts             |  99 +++++++++++++++--
 3 files changed, 275 insertions(+), 11 deletions(-)

diff --git a/src/core/queriesObserver.ts b/src/core/queriesObserver.ts
index 39c3a05ee6..98a456e931 100644
--- a/src/core/queriesObserver.ts
+++ b/src/core/queriesObserver.ts
@@ -1,12 +1,20 @@
 import { difference, getQueryKeyHashFn, replaceAt } from './utils'
 import { notifyManager } from './notifyManager'
-import type { QueryObserverOptions, QueryObserverResult } from './types'
+import type {
+  QueryObserverOptions,
+  QueryObserverResult,
+  RefetchOptions,
+} from './types'
 import type { QueryClient } from './queryClient'
 import { QueryObserver } from './queryObserver'
 import { Subscribable } from './subscribable'
 
 type QueriesObserverListener = (result: QueryObserverResult[]) => void
 
+interface QueriesRefetchOptions extends RefetchOptions {
+  filter: (observer: QueryObserver) => boolean
+}
+
 export class QueriesObserver extends Subscribable<QueriesObserverListener> {
   private client: QueryClient
   private result: QueryObserverResult[]
@@ -57,6 +65,16 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
     return this.result
   }
 
+  refetch(options?: QueriesRefetchOptions): Promise<QueryObserverResult[]> {
+    let observers = this.observers
+
+    if (options?.filter) {
+      observers = this.observers.filter(options.filter)
+    }
+
+    return Promise.all(observers.map(observer => observer.refetch(options)))
+  }
+
   private updateObservers(): void {
     let hasIndexChange = false
 
diff --git a/src/react/tests/useQueries.test.tsx b/src/react/tests/useQueries.test.tsx
index 2a0e6f8257..2cc8947f4c 100644
--- a/src/react/tests/useQueries.test.tsx
+++ b/src/react/tests/useQueries.test.tsx
@@ -1,7 +1,15 @@
+import { fireEvent, waitFor } from '@testing-library/react'
 import React from 'react'
+import { ErrorBoundary } from 'react-error-boundary'
 
-import { queryKey, renderWithClient, sleep } from './utils'
-import { useQueries, QueryClient, UseQueryResult, QueryCache } from '../..'
+import { mockConsoleError, queryKey, renderWithClient, sleep } from './utils'
+import {
+  useQueries,
+  QueryClient,
+  UseQueryResult,
+  QueryCache,
+  useQueryErrorResetBoundary,
+} from '../..'
 
 describe('useQueries', () => {
   const queryCache = new QueryCache()
@@ -30,4 +38,159 @@ describe('useQueries', () => {
     expect(results[1]).toMatchObject([{ data: 1 }, { data: undefined }])
     expect(results[2]).toMatchObject([{ data: 1 }, { data: 2 }])
   })
+
+  it('should render the correct amount of times in suspense mode', async () => {
+    const key1 = queryKey()
+    const key2 = queryKey()
+    const results: UseQueryResult[][] = []
+
+    let renders = 0
+    let count1 = 10
+    let count2 = 20
+
+    function Page() {
+      renders++
+
+      const [stateKey1, setStateKey1] = React.useState(key1)
+
+      const result = useQueries([
+        {
+          queryKey: stateKey1,
+          queryFn: async () => {
+            count1++
+            await sleep(10)
+            return count1
+          },
+        },
+        {
+          queryKey: key2,
+          queryFn: async () => {
+            count2++
+            await sleep(10)
+            return count2
+          },
+          suspense: true,
+        },
+      ])
+
+      results.push(result)
+
+      return (
+        <button aria-label="toggle" onClick={() => setStateKey1(queryKey())} />
+      )
+    }
+
+    const rendered = renderWithClient(
+      queryClient,
+      <React.Suspense fallback="loading">
+        <Page />
+      </React.Suspense>
+    )
+
+    await sleep(50)
+
+    await waitFor(() => rendered.getByLabelText('toggle'))
+    fireEvent.click(rendered.getByLabelText('toggle'))
+
+    await sleep(50)
+
+    expect(renders).toBe(5)
+    expect(results.length).toBe(3)
+
+    // First load
+    expect(results[0]).toMatchObject([
+      { data: 11, status: 'success' },
+      { data: 21, status: 'success' },
+    ])
+
+    // Set state
+    expect(results[1]).toMatchObject([
+      { data: 11, status: 'success' },
+      { data: 21, status: 'success' },
+    ])
+
+    // Second load
+    expect(results[2]).toMatchObject([
+      { data: 12, status: 'success' },
+      { data: 21, status: 'success' },
+    ])
+  })
+
+  it('should retry fetch if the reset error boundary has been reset with global hook', async () => {
+    const key1 = queryKey()
+    const key2 = queryKey()
+
+    let succeed = false
+    const consoleMock = mockConsoleError()
+
+    function Page() {
+      const results = useQueries<string>([
+        {
+          queryKey: key1,
+          queryFn: async () => {
+            await sleep(10)
+            if (!succeed) {
+              throw new Error('Suspense Error Bingo')
+            } else {
+              return 'data1'
+            }
+          },
+          retry: false,
+          suspense: true,
+        },
+        {
+          queryKey: key2,
+          queryFn: () => 'data2',
+          staleTime: Infinity,
+        },
+      ])
+
+      return (
+        <div>
+          <div>data1: {results[0].data}</div>
+          <div>data2: {results[1].data}</div>
+        </div>
+      )
+    }
+
+    function App() {
+      const { reset } = useQueryErrorResetBoundary()
+      return (
+        <ErrorBoundary
+          onReset={reset}
+          fallbackRender={({ resetErrorBoundary }) => (
+            <div>
+              <div>error boundary</div>
+              <button
+                onClick={() => {
+                  resetErrorBoundary()
+                }}
+              >
+                retry
+              </button>
+            </div>
+          )}
+        >
+          <React.Suspense fallback="Loading...">
+            <Page />
+          </React.Suspense>
+        </ErrorBoundary>
+      )
+    }
+
+    const rendered = renderWithClient(queryClient, <App />)
+
+    await waitFor(() => rendered.getByText('Loading...'))
+    await waitFor(() => rendered.getByText('error boundary'))
+    await waitFor(() => rendered.getByText('retry'))
+    fireEvent.click(rendered.getByText('retry'))
+    await waitFor(() => rendered.getByText('error boundary'))
+    await waitFor(() => rendered.getByText('retry'))
+    succeed = true
+    fireEvent.click(rendered.getByText('retry'))
+    await waitFor(() => rendered.getByText('data1: data1'))
+    await waitFor(() => rendered.getByText('data2: data2'))
+
+    consoleMock.mockRestore()
+  })
 })
diff --git a/src/react/useQueries.ts b/src/react/useQueries.ts
index 81774b9ae8..c2b88329aa 100644
--- a/src/react/useQueries.ts
+++ b/src/react/useQueries.ts
@@ -3,31 +3,114 @@ import React from 'react'
 import { notifyManager } from '../core/notifyManager'
 import { QueriesObserver } from '../core/queriesObserver'
 import { useQueryClient } from './QueryClientProvider'
+import { useQueryErrorResetBoundary } from './QueryErrorResetBoundary'
 import { UseQueryOptions, UseQueryResult } from './types'
 
-export function useQueries(queries: UseQueryOptions[]): UseQueryResult[] {
+export function useQueries<TData = unknown, TError = unknown>(
+  queries: UseQueryOptions[]
+): UseQueryResult<TData, TError>[] {
   const queryClient = useQueryClient()
+  const errorResetBoundary = useQueryErrorResetBoundary()
+  const defaultedQueries: UseQueryOptions[] = []
+
+  let someSuspense = false
+  let someUseErrorBoundary = false
+
+  queries.forEach(options => {
+    const defaulted = queryClient.defaultQueryObserverOptions(options)
+
+    // Include callbacks in batch renders
+    if (defaulted.onError) {
+      defaulted.onError = notifyManager.batchCalls(defaulted.onError)
+    }
+
+    if (defaulted.onSuccess) {
+      defaulted.onSuccess = notifyManager.batchCalls(defaulted.onSuccess)
+    }
+
+    if (defaulted.onSettled) {
+      defaulted.onSettled = notifyManager.batchCalls(defaulted.onSettled)
+    }
+
+    if (defaulted.suspense) {
+      someSuspense = true
+    }
+
+    if (defaulted.useErrorBoundary) {
+      someUseErrorBoundary = true
+    }
+
+    defaultedQueries.push(defaulted)
+  })
+
+  if (someSuspense) {
+    defaultedQueries.forEach(options => {
+      // Always set stale time when using suspense to prevent
+      // fetching again when directly re-mounting after suspense
+      if (typeof options.staleTime !== 'number') {
+        options.staleTime = 1000
+      }
+
+      // Prevent retrying failed query if the error boundary has not been reset yet
+      if (!errorResetBoundary.isReset()) {
+        options.retryOnMount = false
+      }
+    })
+  }
 
   // Create queries observer
   const observerRef = React.useRef<QueriesObserver>()
   const observer =
-    observerRef.current || new QueriesObserver(queryClient, queries)
+    observerRef.current || new QueriesObserver(queryClient, defaultedQueries)
   observerRef.current = observer
 
   // Update queries
   if (observer.hasListeners()) {
-    observer.setQueries(queries)
+    observer.setQueries(defaultedQueries)
   }
 
   const [currentResult, setCurrentResult] = React.useState(() =>
     observer.getCurrentResult()
   )
 
+  let someError
+  let someIsLoading = false
+  let someIsError = false
+
+  currentResult.forEach(result => {
+    if (result.isLoading) {
+      someIsLoading = true
+    }
+
+    if (result.isError) {
+      someIsError = true
+    }
+
+    if (result.error) {
+      someError = result.error
+    }
+  })
+
   // Subscribe to the observer
-  React.useEffect(
-    () => observer.subscribe(notifyManager.batchCalls(setCurrentResult)),
-    [observer]
-  )
+  React.useEffect(() => {
+    errorResetBoundary.clearReset()
+    return observer.subscribe(notifyManager.batchCalls(setCurrentResult))
+  }, [observer, errorResetBoundary])
+
+  // Handle suspense
+  if (someSuspense || someUseErrorBoundary) {
+    if (someSuspense && someIsLoading) {
+      errorResetBoundary.clearReset()
+      const unsubscribe = observer.subscribe()
+      throw observer
+        .refetch({ filter: x => x.getCurrentResult().isLoading })
+        .finally(unsubscribe)
+    }
+
+    if (someIsError) {
+      throw someError
+    }
+  }
 
-  return currentResult
+  return currentResult as UseQueryResult<TData, TError>[]
 }