fix: make sure suspense callbacks get correct data

Author: boschni

src/core/queriesObserver.ts

@@ -12,15 +12,20 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
   private result: QueryObserverResult[]
   private queries: QueryObserverOptions[]
   private observers: QueryObserver[]
+  private observersMap: Record<string, QueryObserver>
 
   constructor(client: QueryClient, queries?: QueryObserverOptions[]) {
     super()
 
     this.client = client
-    this.queries = queries || []
+    this.queries = []
     this.result = []
     this.observers = []
-    this.setQueries(this.queries)
+    this.observersMap = {}
+
+    if (queries) {
+      this.setQueries(queries)
+    }
   }
 
   protected onSubscribe(): void {
@@ -59,66 +64,62 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
   }
 
   getOptimisticResult(queries: QueryObserverOptions[]): QueryObserverResult[] {
-    return queries.map((options, i) => {
-      let observer: QueryObserver | undefined = this.observers[i]
-
+    return queries.map(options => {
       const defaultedOptions = this.client.defaultQueryObserverOptions(options)
-
-      if (
-        !observer ||
-        observer.getCurrentQuery().queryHash !== defaultedOptions.queryHash
-      ) {
-        observer = this.observers.find(
-          x => x.getCurrentQuery().queryHash === defaultedOptions.queryHash
-        )
-      }
-
-      if (!observer) {
-        observer = new QueryObserver(this.client, defaultedOptions)
-      }
-
-      return observer.getOptimisticResult(defaultedOptions)
+      return this.getObserver(defaultedOptions).getOptimisticResult(
+        defaultedOptions
+      )
     })
   }
 
+  private getObserver(options: QueryObserverOptions): QueryObserver {
+    const defaultedOptions = this.client.defaultQueryObserverOptions(options)
+    return (
+      this.observersMap[defaultedOptions.queryHash!] ||
+      new QueryObserver(this.client, defaultedOptions)
+    )
+  }
+
   private updateObservers(notifyOptions?: NotifyOptions): void {
     notifyManager.batch(() => {
       let hasIndexChange = false
 
       const prevObservers = this.observers
-      const newObservers = this.queries.map((options, i) => {
-        let observer: QueryObserver | undefined = prevObservers[i]
+      const prevOberversMap = this.observersMap
+
+      const newResult: QueryObserverResult[] = []
+      const newObservers: QueryObserver[] = []
+      const newObserversMap: Record<string, QueryObserver> = {}
 
+      this.queries.forEach((options, i) => {
         const defaultedOptions = this.client.defaultQueryObserverOptions(
           options
         )
+        const queryHash = defaultedOptions.queryHash!
+        const observer = this.getObserver(defaultedOptions)
 
-        if (
-          !observer ||
-          observer.getCurrentQuery().queryHash !== defaultedOptions.queryHash
-        ) {
-          hasIndexChange = true
-          observer = prevObservers.find(
-            x => x.getCurrentQuery().queryHash === defaultedOptions.queryHash
-          )
+        if (prevOberversMap[queryHash]) {
+          observer.setOptions(defaultedOptions, notifyOptions)
         }
 
-        if (observer) {
-          observer.setOptions(defaultedOptions, notifyOptions)
-          return observer
+        if (observer !== prevObservers[i]) {
+          hasIndexChange = true
         }
 
-        return new QueryObserver(this.client, defaultedOptions)
+        newObservers.push(observer)
+        newResult.push(observer.getCurrentResult())
+        newObserversMap[queryHash] = observer
       })
 
       if (prevObservers.length === newObservers.length && !hasIndexChange) {
         return
       }
 
       this.observers = newObservers
-      this.result = newObservers.map(observer => observer.getCurrentResult())
+      this.observersMap = newObserversMap
+      this.result = newResult
 
-      if (!this.listeners.length) {
+      if (!this.hasListeners()) {
         return
       }
 

src/core/queryObserver.ts

@@ -245,6 +245,21 @@ export class QueryObserver<
     return this.fetch(options)
   }
 
+  fetchOptimistic(
+    options: QueryObserverOptions<TQueryFnData, TError, TData, TQueryData>
+  ): Promise<QueryObserverResult<TData, TError>> {
+    const defaultedOptions = this.client.defaultQueryObserverOptions(options)
+
+    const query = this.client
+      .getQueryCache()
+      .build(
+        this.client,
+        defaultedOptions as QueryOptions<TQueryFnData, TError, TQueryData>
+      )
+
+    return query.fetch().then(() => this.createResult(query, defaultedOptions))
+  }
+
   protected fetch(
     fetchOptions?: ObserverFetchOptions
   ): Promise<QueryObserverResult<TData, TError>> {

src/react/tests/suspense.test.tsx

@@ -133,10 +133,18 @@ describe("useQuery's in Suspense mode", () => {
     const successFn = jest.fn()
 
     function Page() {
-      useQuery([key], () => sleep(10), {
-        suspense: true,
-        onSuccess: successFn,
-      })
+      useQuery(
+        [key],
+        async () => {
+          await sleep(10)
+          return key
+        },
+        {
+          suspense: true,
+          select: () => 'selected',
+          onSuccess: successFn,
+        }
+      )
 
       return <>rendered</>
     }
@@ -151,6 +159,7 @@ describe("useQuery's in Suspense mode", () => {
     await waitFor(() => rendered.getByText('rendered'))
 
     expect(successFn).toHaveBeenCalledTimes(1)
+    expect(successFn).toHaveBeenCalledWith('selected')
   })
 
   it('should call every onSuccess handler within a suspense boundary', async () => {

src/react/useBaseQuery.ts

@@ -41,7 +41,7 @@ export function useBaseQuery<TQueryFnData, TError, TData, TQueryData>(
 
   if (defaultedOptions.suspense) {
     // Always set stale time when using suspense to prevent
-    // fetching again when directly re-mounting after suspense
+    // fetching again when directly mounting after suspending
     if (typeof defaultedOptions.staleTime !== 'number') {
       defaultedOptions.staleTime = 1000
     }
@@ -65,9 +65,10 @@ export function useBaseQuery<TQueryFnData, TError, TData, TQueryData>(
   React.useEffect(() => {
     mountedRef.current = true
 
+    errorResetBoundary.clearReset()
+
     const unsubscribe = obsRef.current!.subscribe(
       notifyManager.batchCalls(() => {
-        errorResetBoundary.clearReset()
         if (mountedRef.current) {
           forceUpdate(x => x + 1)
         }
@@ -91,10 +92,10 @@ export function useBaseQuery<TQueryFnData, TError, TData, TQueryData>(
   }, [defaultedOptions])
 
   // Handle suspense
-  if (obsRef.current.options.suspense && result.isLoading) {
-    throw queryClient
-      .fetchQuery(defaultedOptions)
-      .then(data => {
+  if (defaultedOptions.suspense && result.isLoading) {
+    throw obsRef.current
+      .fetchOptimistic(defaultedOptions)
+      .then(({ data }) => {
         defaultedOptions.onSuccess?.(data)
         defaultedOptions.onSettled?.(data, null)
       })
@@ -107,15 +108,14 @@ export function useBaseQuery<TQueryFnData, TError, TData, TQueryData>(
 
   // Handle error boundary
   if (
-    (obsRef.current.options.suspense ||
-      obsRef.current.options.useErrorBoundary) &&
+    (defaultedOptions.suspense || defaultedOptions.useErrorBoundary) &&
     result.isError
   ) {
     throw result.error
   }
 
   // Handle result property usage tracking
-  if (obsRef.current.options.notifyOnChangeProps === 'tracked') {
+  if (defaultedOptions.notifyOnChangeProps === 'tracked') {
     result = obsRef.current.trackResult(result)
   }