fix(useQueries): search for unused observers if keepPreviousData (TanStack#2680) (TanStack#2866)

Author: jvuoti

src/core/queriesObserver.ts

@@ -64,57 +64,94 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
   }
 
   getOptimisticResult(queries: QueryObserverOptions[]): QueryObserverResult[] {
-    return queries.map((options, index) => {
-      const defaultedOptions = this.client.defaultQueryObserverOptions(options)
-      return this.getObserver(defaultedOptions, index).getOptimisticResult(
-        defaultedOptions
-      )
-    })
+    return this.findMatchingObservers(queries).map(match =>
+      match.observer.getCurrentResult()
+    )
   }
 
-  private getObserver(
-    options: QueryObserverOptions,
-    index: number
-  ): QueryObserver {
+  private findMatchingObservers(
+    queries: QueryObserverOptions[]
+  ): QueryObserverMatch[] {
+    const prevObservers = this.observers
+    const defaultedQueryOptions = queries.map(options =>
+      this.client.defaultQueryObserverOptions(options)
+    )
+
+    const matchingObservers: QueryObserverMatch[] = defaultedQueryOptions.flatMap(
+      defaultedOptions => {
+        const match = prevObservers.find(
+          observer => observer.options.queryHash === defaultedOptions.queryHash
+        )
+        if (match != null) {
+          return [{ defaultedQueryOptions: defaultedOptions, observer: match }]
+        }
+        return []
+      }
+    )
+
+    const matchedQueryHashes = matchingObservers.map(
+      match => match.defaultedQueryOptions.queryHash
+    )
+    const unmatchedQueries = defaultedQueryOptions.filter(
+      defaultedOptions =>
+        !matchedQueryHashes.includes(defaultedOptions.queryHash)
+    )
+
+    const unmatchedObservers = prevObservers.filter(
+      prevObserver =>
+        !matchingObservers.some(match => match.observer === prevObserver)
+    )
+
+    const newOrReusedObservers: QueryObserverMatch[] = unmatchedQueries.map(
+      (options, index) => {
+        if (options.keepPreviousData) {
+          // return previous data from one of the observers that no longer match
+          const previouslyUsedObserver = unmatchedObservers[index]
+          if (previouslyUsedObserver !== undefined) {
+            return {
+              defaultedQueryOptions: options,
+              observer: previouslyUsedObserver,
+            }
+          }
+        }
+        return {
+          defaultedQueryOptions: options,
+          observer: this.getObserver(options),
+        }
+      }
+    )
+
+    return matchingObservers.concat(newOrReusedObservers)
+  }
+
+  private getObserver(options: QueryObserverOptions): QueryObserver {
     const defaultedOptions = this.client.defaultQueryObserverOptions(options)
-    let currentObserver = this.observersMap[defaultedOptions.queryHash!]
-    if (!currentObserver && defaultedOptions.keepPreviousData) {
-      currentObserver = this.observers[index]
-    }
+    const currentObserver = this.observersMap[defaultedOptions.queryHash!]
     return currentObserver ?? new QueryObserver(this.client, defaultedOptions)
   }
 
   private updateObservers(notifyOptions?: NotifyOptions): void {
     notifyManager.batch(() => {
-      let hasIndexChange = false
-
       const prevObservers = this.observers
-      const prevObserversMap = this.observersMap
 
-      const newResult: QueryObserverResult[] = []
-      const newObservers: QueryObserver[] = []
-      const newObserversMap: Record<string, QueryObserver> = {}
+      const newObserverMatches = this.findMatchingObservers(this.queries)
 
-      this.queries.forEach((options, i) => {
-        const defaultedOptions = this.client.defaultQueryObserverOptions(
-          options
-        )
-        const queryHash = defaultedOptions.queryHash!
-        const observer = this.getObserver(defaultedOptions, i)
-
-        if (prevObserversMap[queryHash] || defaultedOptions.keepPreviousData) {
-          observer.setOptions(defaultedOptions, notifyOptions)
-        }
-
-        if (observer !== prevObservers[i]) {
-          hasIndexChange = true
-        }
+      // set options for the new observers to notify of changes
+      newObserverMatches.forEach(match =>
+        match.observer.setOptions(match.defaultedQueryOptions, notifyOptions)
+      )
 
-        newObservers.push(observer)
-        newResult.push(observer.getCurrentResult())
-        newObserversMap[queryHash] = observer
-      })
+      const newObservers = newObserverMatches.map(match => match.observer)
+      const newObserversMap = Object.fromEntries(
+        newObservers.map(observer => [observer.options.queryHash, observer])
+      )
+      const newResult = newObservers.map(observer =>
+        observer.getCurrentResult()
+      )
 
+      const hasIndexChange = newObservers.some(
+        (observer, index) => observer !== prevObservers[index]
+      )
       if (prevObservers.length === newObservers.length && !hasIndexChange) {
         return
       }
@@ -157,3 +194,8 @@ export class QueriesObserver extends Subscribable<QueriesObserverListener> {
     })
   }
 }
+
+type QueryObserverMatch = {
+  defaultedQueryOptions: QueryObserverOptions
+  observer: QueryObserver
+}

src/core/tests/queriesObserver.test.tsx

@@ -38,6 +38,51 @@ describe('queriesObserver', () => {
     expect(observerResult).toMatchObject([{ data: 1 }, { data: 2 }])
   })
 
+  test('should still return value for undefined query key', async () => {
+    const key1 = queryKey()
+    const queryFn1 = jest.fn().mockReturnValue(1)
+    const queryFn2 = jest.fn().mockReturnValue(2)
+    const observer = new QueriesObserver(queryClient, [
+      { queryKey: key1, queryFn: queryFn1 },
+      { queryKey: undefined, queryFn: queryFn2 },
+    ])
+    let observerResult
+    const unsubscribe = observer.subscribe(result => {
+      observerResult = result
+    })
+    await sleep(1)
+    unsubscribe()
+    expect(observerResult).toMatchObject([{ data: 1 }, { data: 2 }])
+  })
+
+  test('should return same value for multiple falsy query keys', async () => {
+    const queryFn1 = jest.fn().mockReturnValue(1)
+    const queryFn2 = jest.fn().mockReturnValue(2)
+    const observer = new QueriesObserver(queryClient, [
+      { queryKey: undefined, queryFn: queryFn1 },
+    ])
+    const results: QueryObserverResult[][] = []
+    results.push(observer.getCurrentResult())
+    const unsubscribe = observer.subscribe(result => {
+      results.push(result)
+    })
+    await sleep(1)
+    observer.setQueries([
+      { queryKey: undefined, queryFn: queryFn1 },
+      { queryKey: '', queryFn: queryFn2 },
+    ])
+    await sleep(1)
+    unsubscribe()
+    expect(results.length).toBe(4)
+    expect(results[0]).toMatchObject([{ status: 'idle', data: undefined }])
+    expect(results[1]).toMatchObject([{ status: 'loading', data: undefined }])
+    expect(results[2]).toMatchObject([{ status: 'success', data: 1 }])
+    expect(results[3]).toMatchObject([
+      { status: 'success', data: 1 },
+      { status: 'success', data: 1 },
+    ])
+  })
+
   test('should update when a query updates', async () => {
     const key1 = queryKey()
     const key2 = queryKey()