@tanstack/query-core
Version:
The framework agnostic core that powers TanStack Query
1 lines • 5.27 kB
Source Map (JSON)
{"version":3,"sources":["../../src/streamedQuery.ts"],"sourcesContent":["import { addConsumeAwareSignal, addToEnd } from './utils'\nimport type {\n OmitKeyof,\n QueryFunction,\n QueryFunctionContext,\n QueryKey,\n} from './types'\n\ntype BaseStreamedQueryParams<TQueryFnData, TQueryKey extends QueryKey> = {\n streamFn: (\n context: QueryFunctionContext<TQueryKey>,\n ) => AsyncIterable<TQueryFnData> | Promise<AsyncIterable<TQueryFnData>>\n refetchMode?: 'append' | 'reset' | 'replace'\n}\n\ntype SimpleStreamedQueryParams<\n TQueryFnData,\n TQueryKey extends QueryKey,\n> = BaseStreamedQueryParams<TQueryFnData, TQueryKey> & {\n reducer?: never\n initialValue?: never\n}\n\ntype ReducibleStreamedQueryParams<\n TQueryFnData,\n TData,\n TQueryKey extends QueryKey,\n> = BaseStreamedQueryParams<TQueryFnData, TQueryKey> & {\n reducer: (acc: TData, chunk: TQueryFnData) => TData\n initialValue: TData\n}\n\ntype StreamedQueryParams<TQueryFnData, TData, TQueryKey extends QueryKey> =\n | SimpleStreamedQueryParams<TQueryFnData, TQueryKey>\n | ReducibleStreamedQueryParams<TQueryFnData, TData, TQueryKey>\n\n/**\n * This is a helper function to create a query function that streams data from an AsyncIterable.\n * Data will be an Array of all the chunks received.\n * The query will be in a 'pending' state until the first chunk of data is received, but will go to 'success' after that.\n * The query will stay in fetchStatus 'fetching' until the stream ends.\n * @param queryFn - The function that returns an AsyncIterable to stream data from.\n * @param refetchMode - Defines how re-fetches are handled.\n * Defaults to `'reset'`, erases all data and puts the query back into `pending` state.\n * Set to `'append'` to append new data to the existing data.\n * Set to `'replace'` to write all data to the cache once the stream ends.\n * @param reducer - A function to reduce the streamed chunks into the final data.\n * Defaults to a function that appends chunks to the end of the array.\n * @param initialValue - Initial value to be used while the first chunk is being fetched, and returned if the stream yields no values.\n */\nexport function streamedQuery<\n TQueryFnData = unknown,\n TData = Array<TQueryFnData>,\n TQueryKey extends QueryKey = QueryKey,\n>({\n streamFn,\n refetchMode = 'reset',\n reducer = (items, chunk) =>\n addToEnd(items as Array<TQueryFnData>, chunk) as TData,\n initialValue = [] as TData,\n}: StreamedQueryParams<TQueryFnData, TData, TQueryKey>): QueryFunction<\n TData,\n TQueryKey\n> {\n return async (context) => {\n const query = context.client\n .getQueryCache()\n .find({ queryKey: context.queryKey, exact: true })\n const isRefetch = !!query && query.isFetched()\n if (isRefetch && refetchMode === 'reset') {\n query.setState({\n ...query.resetState,\n fetchStatus: 'fetching',\n })\n }\n\n let result = initialValue\n\n let cancelled: boolean = false as boolean\n const streamFnContext = addConsumeAwareSignal<\n OmitKeyof<typeof context, 'signal'>\n >(\n {\n client: context.client,\n meta: context.meta,\n queryKey: context.queryKey,\n pageParam: context.pageParam,\n direction: context.direction,\n },\n () => context.signal,\n () => (cancelled = true),\n )\n\n const stream = await streamFn(streamFnContext)\n\n const isReplaceRefetch = isRefetch && refetchMode === 'replace'\n\n for await (const chunk of stream) {\n if (cancelled) {\n break\n }\n\n if (isReplaceRefetch) {\n // don't append to the cache directly when replace-refetching\n result = reducer(result, chunk)\n } else {\n context.client.setQueryData<TData>(context.queryKey, (prev) =>\n reducer(prev === undefined ? initialValue : prev, chunk),\n )\n }\n }\n\n // finalize result: replace-refetching needs to write to the cache\n if (isReplaceRefetch && !cancelled) {\n context.client.setQueryData<TData>(context.queryKey, result)\n }\n\n return context.client.getQueryData(context.queryKey) ?? initialValue\n }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,mBAAgD;AAkDzC,SAAS,cAId;AAAA,EACA;AAAA,EACA,cAAc;AAAA,EACd,UAAU,CAAC,OAAO,cAChB,uBAAS,OAA8B,KAAK;AAAA,EAC9C,eAAe,CAAC;AAClB,GAGE;AACA,SAAO,OAAO,YAAY;AACxB,UAAM,QAAQ,QAAQ,OACnB,cAAc,EACd,KAAK,EAAE,UAAU,QAAQ,UAAU,OAAO,KAAK,CAAC;AACnD,UAAM,YAAY,CAAC,CAAC,SAAS,MAAM,UAAU;AAC7C,QAAI,aAAa,gBAAgB,SAAS;AACxC,YAAM,SAAS;AAAA,QACb,GAAG,MAAM;AAAA,QACT,aAAa;AAAA,MACf,CAAC;AAAA,IACH;AAEA,QAAI,SAAS;AAEb,QAAI,YAAqB;AACzB,UAAM,sBAAkB;AAAA,MAGtB;AAAA,QACE,QAAQ,QAAQ;AAAA,QAChB,MAAM,QAAQ;AAAA,QACd,UAAU,QAAQ;AAAA,QAClB,WAAW,QAAQ;AAAA,QACnB,WAAW,QAAQ;AAAA,MACrB;AAAA,MACA,MAAM,QAAQ;AAAA,MACd,MAAO,YAAY;AAAA,IACrB;AAEA,UAAM,SAAS,MAAM,SAAS,eAAe;AAE7C,UAAM,mBAAmB,aAAa,gBAAgB;AAEtD,qBAAiB,SAAS,QAAQ;AAChC,UAAI,WAAW;AACb;AAAA,MACF;AAEA,UAAI,kBAAkB;AAEpB,iBAAS,QAAQ,QAAQ,KAAK;AAAA,MAChC,OAAO;AACL,gBAAQ,OAAO;AAAA,UAAoB,QAAQ;AAAA,UAAU,CAAC,SACpD,QAAQ,SAAS,SAAY,eAAe,MAAM,KAAK;AAAA,QACzD;AAAA,MACF;AAAA,IACF;AAGA,QAAI,oBAAoB,CAAC,WAAW;AAClC,cAAQ,OAAO,aAAoB,QAAQ,UAAU,MAAM;AAAA,IAC7D;AAEA,WAAO,QAAQ,OAAO,aAAa,QAAQ,QAAQ,KAAK;AAAA,EAC1D;AACF;","names":[]}