fix: include config callbacks in batch render (#1036)

Author: boschni

src/core/queryCache.ts

@@ -133,7 +133,9 @@ export class QueryCache {
     predicate?: QueryPredicate,
     options?: QueryPredicateOptions
   ): Query<TResult, TError>[] {
-    if (!options && (predicate === true || typeof predicate === 'undefined')) {
+    const anyKey = predicate === true || typeof predicate === 'undefined'
+
+    if (anyKey && !options) {
       return this.queriesArray
     }
 
@@ -142,20 +144,36 @@ export class QueryCache {
     if (typeof predicate === 'function') {
       predicateFn = predicate as QueryPredicateFn
     } else {
+      const { exact, active, stale } = options || {}
       const resolvedConfig = this.getResolvedQueryConfig(predicate)
 
       predicateFn = query => {
-        if (
-          options &&
-          ((options.exact && query.queryHash !== resolvedConfig.queryHash) ||
-            (typeof options.active === 'boolean' &&
-              query.isActive() !== options.active) ||
-            (typeof options.stale === 'boolean' &&
-              query.isStale() !== options.stale))
-        ) {
+        // Check query key if needed
+        if (!anyKey) {
+          if (exact) {
+            // Check if the query key matches exactly
+            if (query.queryHash !== resolvedConfig.queryHash) {
+              return false
+            }
+          } else {
+            // Check if the query key matches partially
+            if (!deepIncludes(query.queryKey, resolvedConfig.queryKey)) {
+              return false
+            }
+          }
+        }
+
+        // Check active state if needed
+        if (typeof active === 'boolean' && query.isActive() !== active) {
           return false
         }
-        return deepIncludes(query.queryKey, resolvedConfig.queryKey)
+
+        // Check stale state if needed
+        if (typeof stale === 'boolean' && query.isStale() !== stale) {
+          return false
+        }
+
+        return true
       }
     }
 

src/core/queryObserver.ts

@@ -13,6 +13,13 @@ export type UpdateListener<TResult, TError> = (
   result: QueryResult<TResult, TError>
 ) => void
 
+interface NotifyOptions {
+  globalListeners?: boolean
+  listener?: boolean
+  onError?: boolean
+  onSuccess?: boolean
+}
+
 export class QueryObserver<TResult, TError> {
   config: ResolvedQueryConfig<TResult, TError>
 
@@ -150,13 +157,46 @@ export class QueryObserver<TResult, TError> {
     }
   }
 
-  private notify(global?: boolean): void {
+  private notify(options: NotifyOptions): void {
+    const { config, currentResult, currentQuery, listener } = this
+    const { onSuccess, onSettled, onError } = config
+
     notifyManager.batch(() => {
-      notifyManager.schedule(() => {
-        this.listener?.(this.currentResult)
-      })
-      if (global) {
-        this.config.queryCache.notifyGlobalListeners(this.currentQuery)
+      // First trigger the configuration callbacks
+      if (options.onSuccess) {
+        if (onSuccess) {
+          notifyManager.schedule(() => {
+            onSuccess(currentResult.data!)
+          })
+        }
+        if (onSettled) {
+          notifyManager.schedule(() => {
+            onSettled(currentResult.data!, null)
+          })
+        }
+      } else if (options.onError) {
+        if (onError) {
+          notifyManager.schedule(() => {
+            onError(currentResult.error!)
+          })
+        }
+        if (onSettled) {
+          notifyManager.schedule(() => {
+            onSettled(undefined, currentResult.error!)
+          })
+        }
+      }
+
+      // Then trigger the listener
+      if (options.listener && listener) {
+        notifyManager.schedule(() => {
+          listener(currentResult)
+        })
+      }
+
+      // Then the global listeners
+      if (options.globalListeners) {
+        config.queryCache.notifyGlobalListeners(currentQuery)
       }
     })
   }
@@ -180,7 +220,7 @@ export class QueryObserver<TResult, TError> {
       if (!this.isStale) {
         this.isStale = true
         this.updateResult()
-        this.notify(true)
+        this.notify({ listener: true, globalListeners: true })
       }
     }, timeout)
   }
@@ -327,28 +367,30 @@ export class QueryObserver<TResult, TError> {
       this.updateTimers()
     }
 
-    // Trigger callbacks on success or error
-    if (type === 2) {
-      config.onSuccess?.(currentResult.data!)
-      config.onSettled?.(currentResult.data!, null)
-    } else if (type === 3) {
-      config.onError?.(currentResult.error!)
-      config.onSettled?.(undefined, currentResult.error!)
-    }
-
     // Do not notify if the query was invalidated but the stale state did not changed
     if (type === 4 && currentResult.isStale === prevResult.isStale) {
       return
     }
 
+    // Determine which callbacks to trigger
+    const notifyOptions: NotifyOptions = {}
+
+    if (type === 2) {
+      notifyOptions.onSuccess = true
+    } else if (type === 3) {
+      notifyOptions.onError = true
+    }
+
     if (
-      // Always notify on data or error change
+      // Always notify if notifyOnStatusChange is set
+      config.notifyOnStatusChange ||
+      // Otherwise only notify on data or error change
       currentResult.data !== prevResult.data ||
-      currentResult.error !== prevResult.error ||
-      // Maybe notify on other changes
-      config.notifyOnStatusChange
+      currentResult.error !== prevResult.error
     ) {
-      this.notify()
+      notifyOptions.listener = true
     }
+
+    this.notify(notifyOptions)
   }
 }

src/react/tests/useQuery.test.tsx

@@ -1112,6 +1112,45 @@ describe('useQuery', () => {
     expect(renders).toBe(2)
   })
 
+  it('should batch re-renders including hook callbacks', async () => {
+    const key = queryKey()
+
+    let renders = 0
+    let renderedCount = 0
+
+    const queryFn = async () => {
+      await sleep(10)
+      return 'data'
+    }
+
+    function Page() {
+      const [count, setCount] = React.useState(0)
+      useQuery(key, queryFn, {
+        onSuccess: () => {
+          setCount(x => x + 1)
+        },
+      })
+      useQuery(key, queryFn, {
+        onSuccess: () => {
+          setCount(x => x + 1)
+        },
+      })
+      renders++
+      renderedCount = count
+      return null
+    }
+
+    render(<Page />)
+
+    await waitForMs(20)
+
+    // Should be 2 instead of 5
+    expect(renders).toBe(2)
+
+    // Both callbacks should have been executed
+    expect(renderedCount).toBe(2)
+  })
+
   // See https://github.com/tannerlinsley/react-query/issues/170
   it('should start with status idle if enabled is false', async () => {
     const key1 = queryKey()