@@ -9,10 +9,20 @@ import {
99import { create } from 'zustand' ;
1010import { usePolycentric } from '../../lib/polycentric-hooks' ;
1111
12+ /**
13+ * Either a `Query` value or a function that can produce a `Query` from a previous
14+ * query's results.
15+ */
16+ export type QuerySource =
17+ | Query
18+ | ( ( status : QueryStatus | undefined , data : ArrayBuffer | undefined ) => Query ) ;
19+
1220export type QueryRef = {
1321 data : ArrayBuffer | undefined ;
1422 status : QueryStatus ;
1523 error : string | null ;
24+ successfulServers : number ;
25+ pendingServers : number | undefined ;
1626} ;
1727
1828type QueryArgs = {
@@ -36,13 +46,16 @@ type QueryStoreState = {
3646 subscribe : ( key : string , args : QueryArgs ) => void ;
3747 unsubscribe : ( key : string ) => void ;
3848 refresh : ( key : string , args ?: QueryArgs ) => void ;
49+ extend : ( key : string , args : QueryArgs ) => void ;
3950 invalidate : ( key : string , args ?: QueryArgs ) => void ;
4051} ;
4152
4253const EMPTY_ENTRY : QueryRef = Object . freeze ( {
4354 data : undefined ,
4455 status : QueryStatus . Loading ,
4556 error : null ,
57+ successfulServers : 0 ,
58+ pendingServers : undefined ,
4659} ) ;
4760
4861export const useQueryStore = create < QueryStoreState > ( ( set , get ) => {
@@ -53,7 +66,9 @@ export const useQueryStore = create<QueryStoreState>((set, get) => {
5366 if (
5467 merged . data === prev . data &&
5568 merged . status === prev . status &&
56- merged . error === prev . error
69+ merged . error === prev . error &&
70+ merged . successfulServers === prev . successfulServers &&
71+ merged . pendingServers === prev . pendingServers
5772 ) {
5873 return { } ;
5974 }
@@ -68,6 +83,8 @@ export const useQueryStore = create<QueryStoreState>((set, get) => {
6883 updateQueryRef ( key , {
6984 status : QueryStatus . Loading ,
7085 error : null ,
86+ successfulServers : 0 ,
87+ pendingServers : undefined ,
7188 } ) ;
7289
7390 // Request from rs-core
@@ -79,7 +96,12 @@ export const useQueryStore = create<QueryStoreState>((set, get) => {
7996 // Listen for outputs from relevant servers
8097 const sub = observable . subscribe ( {
8198 next ( result ) {
82- updateQueryRef ( key , { data : result . data , status : result . status } ) ;
99+ updateQueryRef ( key , {
100+ data : result . data ,
101+ status : result . status ,
102+ successfulServers : result . successfulServers ,
103+ pendingServers : result . pendingServers ,
104+ } ) ;
83105 } ,
84106 error ( message ) {
85107 console . warn ( `useQuery[${ key } ] error: ${ message } ` ) ;
@@ -139,6 +161,13 @@ export const useQueryStore = create<QueryStoreState>((set, get) => {
139161 sub . dispose = fetch ( key , next ) ;
140162 } ,
141163
164+ extend ( key , args ) {
165+ const sub = get ( ) . subscriptions . get ( key ) ;
166+ if ( ! sub ) return ;
167+ sub . dispose ( ) ;
168+ sub . dispose = fetch ( key , args ) ;
169+ } ,
170+
142171 invalidate ( key , args ) {
143172 const sub = get ( ) . subscriptions . get ( key ) ;
144173 const target = args ?? sub ?. args ;
@@ -151,6 +180,12 @@ export type UseQueryResult = QueryRef & {
151180 isLoading : boolean ;
152181 /** Re-run the fan-out. Cached data stays visible until the new responses arrive. */
153182 refresh : ( ) => void ;
183+ /**
184+ * Re-run the fan-out, but without updating the subscription's query args.
185+ * This allows doing extra queries to add more data while still having refreshes
186+ * re-run the original query.
187+ */
188+ extend : ( ) => void ;
154189 /**
155190 * Drop the rust-side cache for this key, then re-run the fan-out.
156191 * Optional `opts` overrides the original `QueryOpts` for this run
@@ -195,7 +230,9 @@ export function setQueryCache(
195230 if (
196231 merged . data === prev . data &&
197232 merged . status === prev . status &&
198- merged . error === prev . error
233+ merged . error === prev . error &&
234+ merged . successfulServers === prev . successfulServers &&
235+ merged . pendingServers === prev . pendingServers
199236 ) {
200237 return { } ;
201238 }
@@ -206,23 +243,33 @@ export function setQueryCache(
206243}
207244
208245/**
209- * Subscribe to a rust-core query and share its state across every
210- * consumer using the same `queryKey`. The first consumer kicks off
211- * the rust-side fan-out; subsequent consumers refcount onto the same
212- * subscription. `refresh` / `invalidate` re-run the shared fan-out
213- * for every attached consumer. Set `enabled` to `false` to skip the
214- * subscription entirely (the hook still returns cached state if any).
246+ * Subscribe to a rust-core query and share its state across every consumer using
247+ * the same `queryKey`.
248+ * The first consumer kicks off the rust-side fan-out, and subsequent consumers
249+ * refcount onto the same subscription.
250+ * `refresh()` / `invalidate()` re-run the shared fan-out for every attached consumer.
251+ * Set `enabled` to `false` to skip the subscription entirely (the hook still
252+ * returns cached state if any).
253+ *
254+ * Extending/infinite queries can be done by using the "merge" update mode and
255+ * providing a function for the query source.
256+ * Then, calling `extend()` will keep the query key and subscription the same,
257+ * while triggering a new fan-out that will be added to the existing data.
215258 */
216259export function useQuery (
217260 queryKey : QueryKey ,
218- query : Query ,
261+ querySource : QuerySource ,
219262 opts : QueryOpts = { fetchMode : FetchMode . Default } ,
220263 enabled = true ,
221264) : UseQueryResult {
222265 const client = usePolycentric ( ) ;
223266 const cacheKey = queryKey . join ( '\0' ) ;
224267
225268 const entry = useQueryStore ( ( s ) => s . queries . get ( cacheKey ) ?? EMPTY_ENTRY ) ;
269+ const query =
270+ typeof querySource === 'function'
271+ ? querySource ( undefined , undefined )
272+ : querySource ;
226273
227274 // Keep `argsRef` pointing at the freshest call args so the
228275 // imperative handlers below always see them without needing the
@@ -240,6 +287,16 @@ export function useQuery(
240287 return {
241288 ...entry ,
242289 isLoading : enabled && entry . status === QueryStatus . Loading ,
290+ extend : ( ) => {
291+ const query =
292+ typeof querySource === 'function'
293+ ? querySource ( entry . status , entry . data )
294+ : querySource ;
295+
296+ useQueryStore
297+ . getState ( )
298+ . extend ( cacheKey , { client, queryKey, query, opts } ) ;
299+ } ,
243300 refresh : ( ) => useQueryStore . getState ( ) . refresh ( cacheKey , argsRef . current ) ,
244301 invalidate : ( overrideOpts ?: QueryOpts ) =>
245302 useQueryStore
0 commit comments