feat(query-core): add custom reducer support to streamedQuery (#9532)

marcog83 · gobbimar · TkDodo · web-flow · commit 8f24580fdadf · 2025-09-04T11:03:41.000+02:00
* feat(query-core): add custom reducer support to streamedQuery Replace maxChunks parameter with flexible reducer function that delegates data aggregation to consumer code. This provides full control over how streamed chunks are combined into the final data structure. Add support for custom placeholderData that works seamlessly with the reducer function, allowing initialization of complex data types beyond simple arrays. #9065 BREAKING CHANGE: The maxChunks parameter has been removed from streamedQuery. Use a custom reducer function to control data aggregation behavior instead. * ci: apply automated fixes * feat(query-core): require initialValue when using custom reducer in streamedQuery Add type safety by making initialValue mandatory when providing a custom reducer function. This prevents runtime errors and ensures proper data initialization for custom data structures beyond simple arrays. Use conditional types to enforce the relationship between reducer and initialValue parameters, maintaining backward compatibility for simple array-based streaming while requiring explicit initialization for custom reducers. BREAKING CHANGE: When using a custom reducer function with streamedQuery, the initialValue parameter is now required and must be provided. * feat(query-core): require initialValue when using custom reducer in streamedQuery Add type safety by making initialValue mandatory when providing a custom reducer function. This prevents runtime errors and ensures proper data initialization for custom data structures beyond simple arrays. Use conditional types to enforce the relationship between reducer and initialValue parameters, maintaining backward compatibility for simple array-based streaming while requiring explicit initialization for custom reducers. BREAKING CHANGE: When using a custom reducer function with streamedQuery, the initialValue parameter is now required and must be provided. * feat(query-core): require initialValue when using custom reducer in streamedQuery Add type safety by making initialValue mandatory when providing a custom reducer function. This prevents runtime errors and ensures proper data initialization for custom data structures beyond simple arrays. Use conditional types to enforce the relationship between reducer and initialValue parameters, maintaining backward compatibility for simple array-based streaming while requiring explicit initialization for custom reducers. BREAKING CHANGE: When using a custom reducer function with streamedQuery, the initialValue parameter is now required and must be provided. * removed personal vscode workspace file * updated documentation * fix(docs): clarify reducer function description in streamedQuery documentation * ci: apply automated fixes * fix(tests): Code Review :: update streamedQuery tests to use correct initialValue type --------- Co-authored-by: gobbimar <marco.gobbi@adidas.com> Co-authored-by: Dominik Dorfmeister <office@dorfmeister.cc> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
diff --git a/docs/reference/streamedQuery.md b/docs/reference/streamedQuery.md
@@ -33,9 +33,13 @@ const query = queryOptions({
   - When set to `'reset'`, the query will erase all data and go back into `pending` state.
   - When set to `'append'`, data will be appended to existing data.
   - When set to `'replace'`, all data will be written to the cache once the stream ends.
-- `maxChunks?: number`
+- `reducer?: (accumulator: TData, chunk: TQueryFnData) => TData`
   - Optional
-  - The maximum number of chunks to keep in the cache.
-  - Defaults to `undefined`, meaning all chunks will be kept.
-  - If `undefined` or `0`, the number of chunks is unlimited.
-  - If the number of chunks exceeds this number, the oldest chunk will be removed.
+  - Reduces streamed chunks (`TQueryFnData`) into the final data shape (`TData`).
+  - Default: appends each chunk to the end of the accumulator when `TData` is an array.
+  - If `TData` is not an array, you must provide a custom `reducer`.
+- `initialValue?: TData = TQueryFnData`
+  - Optional
+  - Defines the initial data to be used while the first chunk is being fetched.
+  - It is mandatory when custom `reducer` is provided.
+  - Defaults to an empty array.
diff --git a/packages/query-core/src/__tests__/streamedQuery.test.tsx b/packages/query-core/src/__tests__/streamedQuery.test.tsx
@@ -391,13 +391,18 @@ describe('streamedQuery', () => {
     })
   })
 
-  test('should support maxChunks', async () => {
+  test('should support custom reducer', async () => {
     const key = queryKey()
+
     const observer = new QueryObserver(queryClient, {
       queryKey: key,
       queryFn: streamedQuery({
-        queryFn: () => createAsyncNumberGenerator(3),
-        maxChunks: 2,
+        queryFn: () => createAsyncNumberGenerator(2),
+        reducer: (acc, chunk) => ({
+          ...acc,
+          [chunk]: true,
+        }),
+        initialValue: {} as Record<number, boolean>,
       }),
     })
 
@@ -409,41 +414,34 @@ describe('streamedQuery', () => {
       data: undefined,
     })
 
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [0],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [0, 1],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
+    await vi.advanceTimersByTimeAsync(100)
 
     expect(observer.getCurrentResult()).toMatchObject({
       status: 'success',
       fetchStatus: 'idle',
-      data: [1, 2],
+      data: {
+        0: true,
+        1: true,
+      },
     })
 
     unsubscribe()
   })
 
-  test('maxChunks with append refetch', async () => {
+  test('should support custom reducer with initialValue', async () => {
     const key = queryKey()
     const observer = new QueryObserver(queryClient, {
       queryKey: key,
       queryFn: streamedQuery({
-        queryFn: () => createAsyncNumberGenerator(3),
-        maxChunks: 2,
-        refetchMode: 'append',
+        queryFn: () => createAsyncNumberGenerator(2),
+        reducer: (acc, chunk) => ({
+          ...acc,
+          [chunk]: true,
+        }),
+        initialValue: {
+          10: true,
+          11: true,
+        } as Record<number, boolean>,
       }),
     })
 
@@ -455,62 +453,17 @@ describe('streamedQuery', () => {
       data: undefined,
     })
 
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [0],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [0, 1],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'idle',
-      data: [1, 2],
-    })
-
-    void observer.refetch()
-
-    await vi.advanceTimersByTimeAsync(10)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [1, 2],
-    })
-
-    await vi.advanceTimersByTimeAsync(40)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [2, 0],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
-
-    expect(observer.getCurrentResult()).toMatchObject({
-      status: 'success',
-      fetchStatus: 'fetching',
-      data: [0, 1],
-    })
-
-    await vi.advanceTimersByTimeAsync(50)
+    await vi.advanceTimersByTimeAsync(100)
 
     expect(observer.getCurrentResult()).toMatchObject({
       status: 'success',
       fetchStatus: 'idle',
-      data: [1, 2],
+      data: {
+        10: true,
+        11: true,
+        0: true,
+        1: true,
+      },
     })
 
     unsubscribe()
diff --git a/packages/query-core/src/streamedQuery.ts b/packages/query-core/src/streamedQuery.ts
@@ -1,6 +1,34 @@
 import { addToEnd } from './utils'
 import type { QueryFunction, QueryFunctionContext, QueryKey } from './types'
 
+type BaseStreamedQueryParams<TQueryFnData, TQueryKey extends QueryKey> = {
+  queryFn: (
+    context: QueryFunctionContext<TQueryKey>,
+  ) => AsyncIterable<TQueryFnData> | Promise<AsyncIterable<TQueryFnData>>
+  refetchMode?: 'append' | 'reset' | 'replace'
+}
+
+type SimpleStreamedQueryParams<
+  TQueryFnData,
+  TQueryKey extends QueryKey,
+> = BaseStreamedQueryParams<TQueryFnData, TQueryKey> & {
+  reducer?: never
+  initialValue?: never
+}
+
+type ReducibleStreamedQueryParams<
+  TQueryFnData,
+  TData,
+  TQueryKey extends QueryKey,
+> = BaseStreamedQueryParams<TQueryFnData, TQueryKey> & {
+  reducer: (acc: TData, chunk: TQueryFnData) => TData
+  initialValue: TData
+}
+
+type StreamedQueryParams<TQueryFnData, TData, TQueryKey extends QueryKey> =
+  | SimpleStreamedQueryParams<TQueryFnData, TQueryKey>
+  | ReducibleStreamedQueryParams<TQueryFnData, TData, TQueryKey>
+
 /**
  * This is a helper function to create a query function that streams data from an AsyncIterable.
  * Data will be an Array of all the chunks received.
@@ -11,31 +39,29 @@ import type { QueryFunction, QueryFunctionContext, QueryKey } from './types'
  * Defaults to `'reset'`, erases all data and puts the query back into `pending` state.
  * Set to `'append'` to append new data to the existing data.
  * Set to `'replace'` to write all data to the cache once the stream ends.
- * @param maxChunks - The maximum number of chunks to keep in the cache.
- * Defaults to `undefined`, meaning all chunks will be kept.
- * If `undefined` or `0`, the number of chunks is unlimited.
- * If the number of chunks exceeds this number, the oldest chunk will be removed.
+ * @param reducer - A function to reduce the streamed chunks into the final data.
+ * Defaults to a function that appends chunks to the end of the array.
+ * @param initialValue - Initial value to be used while the first chunk is being fetched.
  */
 export function streamedQuery<
   TQueryFnData = unknown,
+  TData = Array<TQueryFnData>,
   TQueryKey extends QueryKey = QueryKey,
 >({
   queryFn,
   refetchMode = 'reset',
-  maxChunks,
-}: {
-  queryFn: (
-    context: QueryFunctionContext<TQueryKey>,
-  ) => AsyncIterable<TQueryFnData> | Promise<AsyncIterable<TQueryFnData>>
-  refetchMode?: 'append' | 'reset' | 'replace'
-  maxChunks?: number
-}): QueryFunction<Array<TQueryFnData>, TQueryKey> {
+  reducer = (items, chunk) =>
+    addToEnd(items as Array<TQueryFnData>, chunk) as TData,
+  initialValue = [] as TData,
+}: StreamedQueryParams<TQueryFnData, TData, TQueryKey>): QueryFunction<
+  TData,
+  TQueryKey
+> {
   return async (context) => {
     const query = context.client
       .getQueryCache()
       .find({ queryKey: context.queryKey, exact: true })
     const isRefetch = !!query && query.state.data !== undefined
-
     if (isRefetch && refetchMode === 'reset') {
       query.setState({
         status: 'pending',
@@ -45,7 +71,8 @@ export function streamedQuery<
       })
     }
 
-    let result: Array<TQueryFnData> = []
+    let result = initialValue
+
     const stream = await queryFn(context)
 
     for await (const chunk of stream) {
@@ -55,19 +82,16 @@ export function streamedQuery<
 
       // don't append to the cache directly when replace-refetching
       if (!isRefetch || refetchMode !== 'replace') {
-        context.client.setQueryData<Array<TQueryFnData>>(
-          context.queryKey,
-          (prev = []) => {
-            return addToEnd(prev, chunk, maxChunks)
-          },
+        context.client.setQueryData<TData>(context.queryKey, (prev) =>
+          reducer(prev === undefined ? initialValue : prev, chunk),
         )
       }
-      result = addToEnd(result, chunk, maxChunks)
+      result = reducer(result, chunk)
     }
 
     // finalize result: replace-refetching needs to write to the cache
     if (isRefetch && refetchMode === 'replace' && !context.signal.aborted) {
-      context.client.setQueryData<Array<TQueryFnData>>(context.queryKey, result)
+      context.client.setQueryData<TData>(context.queryKey, result)
     }
 
     return context.client.getQueryData(context.queryKey)!
