File size: 3,981 Bytes
74bd15c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | import { addConsumeAwareSignal, addToEnd } from './utils'
import type {
OmitKeyof,
QueryFunction,
QueryFunctionContext,
QueryKey,
} from './types'
type BaseStreamedQueryParams<TQueryFnData, TQueryKey extends QueryKey> = {
streamFn: (
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.
* The query will be in a 'pending' state until the first chunk of data is received, but will go to 'success' after that.
* The query will stay in fetchStatus 'fetching' until the stream ends.
* @param queryFn - The function that returns an AsyncIterable to stream data from.
* @param refetchMode - Defines how re-fetches are handled.
* 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 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, and returned if the stream yields no values.
*/
export function streamedQuery<
TQueryFnData = unknown,
TData = Array<TQueryFnData>,
TQueryKey extends QueryKey = QueryKey,
>({
streamFn,
refetchMode = 'reset',
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.isFetched()
if (isRefetch && refetchMode === 'reset') {
query.setState({
...query.resetState,
fetchStatus: 'fetching',
})
}
let result = initialValue
let cancelled: boolean = false as boolean
const streamFnContext = addConsumeAwareSignal<
OmitKeyof<typeof context, 'signal'>
>(
{
client: context.client,
meta: context.meta,
queryKey: context.queryKey,
pageParam: context.pageParam,
direction: context.direction,
},
() => context.signal,
() => (cancelled = true),
)
const stream = await streamFn(streamFnContext)
const isReplaceRefetch = isRefetch && refetchMode === 'replace'
for await (const chunk of stream) {
if (cancelled) {
break
}
if (isReplaceRefetch) {
// don't append to the cache directly when replace-refetching
result = reducer(result, chunk)
} else {
context.client.setQueryData<TData>(context.queryKey, (prev) =>
reducer(prev === undefined ? initialValue : prev, chunk),
)
}
}
// finalize result: replace-refetching needs to write to the cache
if (isReplaceRefetch && !cancelled) {
context.client.setQueryData<TData>(context.queryKey, result)
}
return context.client.getQueryData(context.queryKey) ?? initialValue
}
}
|