feat: add useQueryState hook

Author: crutchcorn

packages/react-query/src/index.ts

@@ -40,6 +40,7 @@ export {
 } from './QueryErrorResetBoundary'
 export { useIsFetching } from './useIsFetching'
 export { useIsMutating, useMutationState } from './useMutationState'
+export { useQueryState } from './useQueryState'
 export { useMutation } from './useMutation'
 export { useInfiniteQuery } from './useInfiniteQuery'
 export { useIsRestoring, IsRestoringProvider } from './isRestoring'

packages/react-query/src/useIsFetching.ts

@@ -1,24 +1,13 @@
 'use client'
-import * as React from 'react'
-import { notifyManager } from '@tanstack/query-core'
-
-import { useQueryClient } from './QueryClientProvider'
+import { useQueryState } from './useQueryState'
 import type { QueryClient, QueryFilters } from '@tanstack/query-core'
 
 export function useIsFetching(
   filters?: QueryFilters,
   queryClient?: QueryClient,
 ): number {
-  const client = useQueryClient(queryClient)
-  const queryCache = client.getQueryCache()
-
-  return React.useSyncExternalStore(
-    React.useCallback(
-      (onStoreChange) =>
-        queryCache.subscribe(notifyManager.batchCalls(onStoreChange)),
-      [queryCache],
-    ),
-    () => client.isFetching(filters),
-    () => client.isFetching(filters),
-  )
+  return useQueryState(
+    { filters: { ...filters, fetchStatus: 'fetching' } },
+    queryClient,
+  ).length
 }

packages/react-query/src/useQueryState.ts

@@ -0,0 +1,70 @@
+'use client'
+import * as React from 'react'
+
+import { notifyManager, replaceEqualDeep } from '@tanstack/query-core'
+import { useQueryClient } from './QueryClientProvider'
+import type {
+  DefaultError,
+  Query,
+  QueryCache,
+  QueryClient,
+  QueryFilters,
+  QueryKey,
+  QueryState,
+} from '@tanstack/query-core'
+
+type QueryStateOptions<TResult = QueryState> = {
+  filters?: QueryFilters
+  select?: (query: Query<unknown, DefaultError, unknown, QueryKey>) => TResult
+}
+
+function getResult<TResult = QueryState>(
+  queryCache: QueryCache,
+  options: QueryStateOptions<TResult>,
+): Array<TResult> {
+  return queryCache
+    .findAll(options.filters)
+    .map(
+      (query): TResult =>
+        (options.select
+          ? options.select(
+              query as Query<unknown, DefaultError, unknown, QueryKey>,
+            )
+          : query.state) as TResult,
+    )
+}
+
+export function useQueryState<TResult = QueryState>(
+  options: QueryStateOptions<TResult> = {},
+  queryClient?: QueryClient,
+): Array<TResult> {
+  const queryCache = useQueryClient(queryClient).getQueryCache()
+  const optionsRef = React.useRef(options)
+  const result = React.useRef<Array<TResult>>()
+  if (!result.current) {
+    result.current = getResult(queryCache, options)
+  }
+
+  React.useEffect(() => {
+    optionsRef.current = options
+  })
+
+  return React.useSyncExternalStore(
+    React.useCallback(
+      (onStoreChange) =>
+        queryCache.subscribe(() => {
+          const nextResult = replaceEqualDeep(
+            result.current,
+            getResult(queryCache, optionsRef.current),
+          )
+          if (result.current !== nextResult) {
+            result.current = nextResult
+            notifyManager.schedule(onStoreChange)
+          }
+        }),
+      [queryCache],
+    ),
+    () => result.current,
+    () => result.current,
+  )!
+}