1import { QueryClient, type QueryFunction, type QueryKey, type Updater } from "@tanstack/react-query";
2import type { OptimisticEvents, Reactive } from "./client.ts";
3import { type AllObjectPaths, type GetObjectPath, getPath, setPath } from "./object-path.ts";
4
5export type QueryKeyAndFn<T = unknown, Key extends QueryKey = QueryKey> = {
6 queryKey: Key;
7 queryFn?: QueryFunction<T, any, never> | undefined;
8};
9
10type RequireAtLeastOneKey<T> = {
11 [K in keyof T]-?: Required<Pick<T, K>> & Partial<Omit<T, K>>;
12}[keyof T];
13
14export function boundQueryClientGet(
15 queryClient: QueryClient,
16): <T>({ queryKey }: QueryKeyAndFn<T>) => T | undefined {
17 return function get<T>({ queryKey }: QueryKeyAndFn<T>) {
18 return queryClient.getQueryData<T>(queryKey);
19 };
20}
21
22export function reactiveFromQueryCache<T>(
23 client: QueryClient,
24 { queryKey }: QueryKeyAndFn<T>,
25): Reactive<T | undefined>;
26export function reactiveFromQueryCache<T, R>(
27 client: QueryClient,
28 { queryKey }: QueryKeyAndFn<T>,
29 deriver: (value: T | undefined) => R,
30): Reactive<R>;
31export function reactiveFromQueryCache<T>(
32 client: QueryClient,
33 { queryKey }: QueryKeyAndFn<T>,
34 deriver: (value: T | undefined) => unknown = x => x,
35): Reactive<unknown> {
36 return {
37 get: () => deriver(client.getQueryData<T>(queryKey)),
38 sub: (onChange) => {
39 const queryKeyJson = JSON.stringify(queryKey);
40 return client.getQueryCache().subscribe((event) => {
41 if (JSON.stringify(event.query.queryKey) === queryKeyJson) {
42 onChange();
43 }
44 });
45 },
46 };
47}
48
49class TanstackQueryOptimisticHelpers {
50 #client: QueryClient;
51 #onRefetch: OptimisticEvents["onRefetch"];
52 #onRestore: OptimisticEvents["onRestore"];
53 #refetchHashes: string[] = [];
54
55 constructor(client: QueryClient, { onRefetch, onRestore }: OptimisticEvents) {
56 this.#client = client;
57 this.#onRefetch = onRefetch;
58 this.#onRestore = onRestore;
59 }
60
61 #get<T>({ queryKey }: QueryKeyAndFn<T>) {
62 return this.#client.getQueryData<T>(queryKey);
63 }
64
65 #set<T>(
66 query: QueryKeyAndFn<T>,
67 updater: Updater<NoInfer<T> | undefined, NoInfer<T> | undefined>,
68 ) {
69 const result = this.#client.setQueryData<T>(query.queryKey, updater);
70 if (result) {
71 this.#client.cancelQueries({ queryKey: query.queryKey, exact: true })
72 .catch(() => {});
73 this.refetchOnSettled(query);
74 }
75 }
76
77 /**
78 * Set the entire query data.
79 * If the query doesn't exist, the new query is created.
80 */
81 set = <Data>(
82 queryKey: QueryKeyAndFn<Data>,
83 value: Data | ((prev: Data | undefined) => Data | undefined),
84 ): void => {
85 const prev = this.#get(queryKey);
86
87 const newValue = typeof value === "function"
88 ? (value as (prev: Data | undefined) => Data | undefined)(prev)
89 : value;
90
91 this.#set(queryKey, newValue);
92 this.#onRestore(() => {
93 if (prev === undefined) {
94 this.#client.removeQueries({
95 queryKey: queryKey.queryKey,
96 exact: true,
97 });
98 } else {
99 this.#set(queryKey, prev);
100 }
101 });
102 };
103
104 /**
105 * Update the entire query data.
106 * If the query doesn't exist, it cancels
107 */
108 updateExisting = <Data>(
109 queryKey: QueryKeyAndFn<Data>,
110 value: Data | ((prev: Data) => Data),
111 ) => {
112 const prev = this.#get(queryKey);
113 if (!prev) return;
114
115 const newValue = typeof value === "function"
116 ? (value as (prev: Data | undefined) => Data | undefined)(prev)
117 : value;
118
119 this.#set(queryKey, newValue);
120 this.#onRestore(() => {
121 this.#set(queryKey, prev);
122 });
123 };
124
125 /**
126 * Mark a query as changed and schedule a refetch.
127 * Does not modify any data.
128 * If the query doesn't exist, the updater is skipped.
129 */
130 refetchOnSettled = (queryKey: QueryKeyAndFn) => {
131 if (!queryKey.queryFn) return;
132 const state = this.#client.getQueryCache().find({
133 queryKey: queryKey.queryKey,
134 exact: true,
135 });
136 if (!state || this.#refetchHashes.includes(state.queryHash)) return;
137 this.#refetchHashes.push(state.queryHash);
138 try {
139 this.#onRefetch(async () => {
140 await this.#client.invalidateQueries(queryKey);
141 });
142 } catch { /* Ignore expiry */ }
143 };
144
145 /**
146 * Set a property at an object path.
147 * If the query or path doesn't exist, the updater is skipped.
148 */
149 objSet = <Data extends object, const Path extends AllObjectPaths<Data>>(
150 queryKey: QueryKeyAndFn<Data>,
151 path: Path,
152 value:
153 | Exclude<GetObjectPath<Data, Path>, Function>
154 | ((prev: GetObjectPath<Data, Path>) => GetObjectPath<Data, Path>),
155 ) => {
156 const prev = this.#get(queryKey);
157 if (!prev) return;
158 const { value: original, exists } = getPath(prev, path);
159 if (!exists) return;
160
161 this.#set(
162 queryKey,
163 (obj) =>
164 obj
165 ? setPath(
166 obj,
167 path,
168 typeof value === "function"
169 ? (value as ((
170 prev: GetObjectPath<Data, Path>,
171 ) => GetObjectPath<Data, Path>))(original)
172 : value,
173 )
174 : obj,
175 );
176 this.#onRestore(() => {
177 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
178 });
179 };
180
181 /**
182 * Increment a numeric property at an object path.
183 * If the query or path doesn't exist, the updater is skipped.
184 */
185 objIncrement = <
186 Data extends object,
187 const Path extends AllObjectPaths<Data>,
188 >(
189 queryKey: QueryKeyAndFn<Data>,
190 path: Path,
191 amount: number = 1,
192 ) => {
193 const prev = this.#get(queryKey);
194 if (!prev) return;
195 const { value: original, exists } = getPath(prev, path);
196 if (!exists || typeof original !== "number") return;
197
198 this.#set(
199 queryKey,
200 (obj) => obj ? setPath(obj, path, (original + amount) as any) : obj,
201 );
202 this.#onRestore(() => {
203 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
204 });
205 };
206
207 /**
208 * Decrement a numeric property at an object path.
209 * If the query or path doesn't exist, the updater is skipped.
210 */
211 objDecrement = <
212 Data extends object,
213 const Path extends AllObjectPaths<Data>,
214 >(
215 queryKey: QueryKeyAndFn<Data>,
216 path: Path,
217 amount: number = 1,
218 ) => {
219 const prev = this.#get(queryKey);
220 if (!prev) return;
221 const { value: original, exists } = getPath(prev, path);
222 if (!exists || typeof original !== "number") return;
223
224 this.#set(
225 queryKey,
226 (obj) => obj ? setPath(obj, path, (original - amount) as any) : obj,
227 );
228 this.#onRestore(() => {
229 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
230 });
231 };
232
233 /**
234 * Toggle a boolean property at an object path.
235 * If the query or path doesn't exist, the updater is skipped.
236 */
237 objToggle = <Data extends object, const Path extends AllObjectPaths<Data>>(
238 queryKey: QueryKeyAndFn<Data>,
239 path: Path,
240 ) => {
241 const prev = this.#get(queryKey);
242 if (!prev) return;
243 const { value: original, exists } = getPath(prev, path);
244 if (!exists || typeof original !== "boolean") return;
245
246 this.#set(
247 queryKey,
248 (obj) => obj ? setPath(obj, path, (!original) as any) : obj,
249 );
250 this.#onRestore(() => {
251 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
252 });
253 };
254
255 /**
256 * Shallow merge multiple properties at an object path.
257 * If the query or path doesn't exist, the updater is skipped.
258 */
259 objSetMany = <Data extends object, const Path extends AllObjectPaths<Data>>(
260 queryKey: QueryKeyAndFn<Data>,
261 path: Path,
262 updates: Partial<GetObjectPath<Data, Path>>,
263 ) => {
264 const prev = this.#get(queryKey);
265 if (!prev) return;
266 const { value: original, exists } = getPath(prev, path);
267 if (!exists || typeof original !== "object" || original === null) {
268 return;
269 }
270
271 const merged = { ...original as object, ...updates } as GetObjectPath<
272 Data,
273 Path
274 >;
275 this.#set(
276 queryKey,
277 (obj) => obj ? setPath(obj, path, merged) : obj,
278 );
279 this.#onRestore(() => {
280 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
281 });
282 };
283
284 /**
285 * Push item(s) to the end of an array at an object path.
286 * If the query or path doesn't exist, the updater is skipped.
287 */
288 objArrayPush = <
289 Data extends object,
290 const Path extends AllObjectPaths<Data>,
291 >(
292 queryKey: QueryKeyAndFn<Data>,
293 path: Path,
294 ...items: GetObjectPath<Data, Path> extends Array<infer T> ? T[]
295 : never
296 ) => {
297 const prev = this.#get(queryKey);
298 if (!prev) return;
299 const { value: original, exists } = getPath(prev, path);
300 if (!exists || !Array.isArray(original)) return;
301
302 const newArray = [...original, ...items];
303 this.#set(
304 queryKey,
305 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
306 );
307 this.#onRestore(() => {
308 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
309 });
310 };
311
312 /**
313 * Add item(s) to the beginning of an array at an object path.
314 * If the query or path doesn't exist, the updater is skipped.
315 */
316 objArrayUnshift = <
317 Data extends object,
318 const Path extends AllObjectPaths<Data>,
319 >(
320 queryKey: QueryKeyAndFn<Data>,
321 path: Path,
322 ...items: GetObjectPath<Data, Path> extends Array<infer T> ? T[]
323 : never
324 ) => {
325 const prev = this.#get(queryKey);
326 if (!prev) return;
327 const { value: original, exists } = getPath(prev, path);
328 if (!exists || !Array.isArray(original)) return;
329
330 const newArray = [...items, ...original];
331 this.#set(
332 queryKey,
333 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
334 );
335 this.#onRestore(() => {
336 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
337 });
338 };
339
340 /**
341 * Remove items from an array that match `filter`.
342 * If the query or path doesn't exist, the updater is skipped.
343 */
344 objArrayRemove = <
345 Data extends object,
346 const Path extends AllObjectPaths<Data>,
347 >(
348 queryKey: QueryKeyAndFn<Data>,
349 path: Path,
350 removeFilter: (
351 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
352 index: number,
353 ) => boolean,
354 ) => {
355 const prev = this.#get(queryKey);
356 if (!prev) return;
357 const { value: original, exists } = getPath(prev, path);
358 if (!exists || !Array.isArray(original)) return;
359
360 const newArray = original.filter((item, index) => !removeFilter(item, index));
361 this.#set(
362 queryKey,
363 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
364 );
365 this.#onRestore(() => {
366 // TODO: splice items back in case original changed
367 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
368 });
369 };
370
371 /**
372 * Filter items to just include items that match `filter`. This is the inverse of `objArrayRemove`
373 * If the query or path doesn't exist, the updater is skipped.
374 */
375 objArrayFilter = <
376 Data extends object,
377 const Path extends AllObjectPaths<Data>,
378 >(
379 queryKey: QueryKeyAndFn<Data>,
380 path: Path,
381 filter: (
382 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
383 index: number,
384 ) => boolean,
385 ) => {
386 const prev = this.#get(queryKey);
387 if (!prev) return;
388 const { value: original, exists } = getPath(prev, path);
389 if (!exists || !Array.isArray(original)) return;
390
391 const newArray = original.filter(filter);
392 this.#set(
393 queryKey,
394 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
395 );
396 this.#onRestore(() => {
397 // TODO: splice items back in case original changed
398 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
399 });
400 };
401
402 /**
403 * Update items in an array that match a predicate.
404 * If the query or path doesn't exist, the updater is skipped.
405 */
406 objArrayUpdate = <
407 Data extends object,
408 const Path extends AllObjectPaths<Data>,
409 >(
410 queryKey: QueryKeyAndFn<Data>,
411 path: Path,
412 {
413 filter,
414 update,
415 }: {
416 filter?: (
417 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
418 index: number,
419 ) => boolean;
420 update: (
421 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
422 ) => GetObjectPath<Data, Path> extends Array<infer T> ? T : never;
423 },
424 ): { inserted: boolean } => {
425 return this.objArrayUpsert(queryKey, path, {
426 filter: filter ?? (() => true),
427 update,
428 });
429 };
430
431 /**
432 * Update items in an array that match a `filter`, or insert a new one if there was no match.
433 * If the query or path doesn't exist, the updater is skipped.
434 */
435 objArrayUpsert = <
436 Data extends object,
437 const Path extends AllObjectPaths<Data>,
438 >(
439 queryKey: QueryKeyAndFn<Data>,
440 path: Path,
441 {
442 filter,
443 update,
444 insert,
445 }:
446 & {
447 filter: (
448 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
449 index: number,
450 ) => boolean;
451 }
452 & RequireAtLeastOneKey<{
453 /** Defaults to the identity function */
454 update: (
455 item: GetObjectPath<Data, Path> extends Array<infer T> ? T : never,
456 ) => GetObjectPath<Data, Path> extends Array<infer T> ? T : never;
457 /** Defaults to not inserting */
458 insert: () => GetObjectPath<Data, Path> extends Array<infer T> ? T : never;
459 }>,
460 ): { inserted: boolean } => {
461 const prev = this.#get(queryKey);
462 if (!prev) return { inserted: false };
463 const { value: original, exists } = getPath(prev, path);
464 if (!exists || !Array.isArray(original)) return { inserted: false };
465
466 let matched = false;
467 const newArray = original.map((item, index) => {
468 if (filter(item, index)) {
469 matched = true;
470 return update ? update(item) : item;
471 } else {
472 return item;
473 }
474 });
475 let inserted = false;
476 if (!matched && insert) {
477 newArray.push(insert());
478 inserted = true;
479 }
480 this.#set(
481 queryKey,
482 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
483 );
484 this.#onRestore(() => {
485 // TODO: splice items back in case original changed
486 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
487 });
488 return { inserted };
489 };
490
491 /**
492 * Insert item(s) at a specific index in an array.
493 * If the query or path doesn't exist, the updater is skipped.
494 */
495 objArrayInsertIndex = <
496 Data extends object,
497 const Path extends AllObjectPaths<Data>,
498 >(
499 queryKey: QueryKeyAndFn<Data>,
500 path: Path,
501 index: number,
502 ...items: GetObjectPath<Data, Path> extends Array<infer T> ? T[]
503 : never
504 ) => {
505 const prev = this.#get(queryKey);
506 if (!prev) return;
507 const { value: original, exists } = getPath(prev, path);
508 if (!exists || !Array.isArray(original)) return;
509
510 const newArray = [
511 ...original.slice(0, index),
512 ...items,
513 ...original.slice(index),
514 ];
515 this.#set(
516 queryKey,
517 (obj) => obj ? setPath(obj, path, newArray as any) : obj,
518 );
519 this.#onRestore(() => {
520 // TODO: splice items back in case original changed
521 this.#set(queryKey, (obj) => obj ? setPath(obj, path, original) : obj);
522 });
523 };
524
525 /**
526 * Push item(s) to the end of an array.
527 * If the query or path doesn't exist, the updater is skipped.
528 */
529 arrayPush = <Data>(queryKey: QueryKeyAndFn<Data[] | null>, ...items: Data[]) => {
530 const prev = this.#get(queryKey);
531 if (!prev || !Array.isArray(prev)) return;
532
533 const newArray = [...prev, ...items];
534 this.#set(queryKey, newArray);
535 this.#onRestore(() => {
536 this.#set(
537 queryKey,
538 (old) => old ? old.filter((x) => !items.includes(x)) : old,
539 );
540 });
541 };
542
543 /**
544 * Add item(s) to the beginning of an array at an object path.
545 * If the query or path doesn't exist, the updater is skipped.
546 */
547 arrayUnshift = <Data>(queryKey: QueryKeyAndFn<Data[] | null>, ...items: Data[]) => {
548 const prev = this.#get(queryKey);
549 if (!prev || !Array.isArray(prev)) return;
550
551 const newArray = [...items, ...prev];
552 this.#set(queryKey, newArray);
553 this.#onRestore(() => {
554 this.#set(
555 queryKey,
556 (old) => old ? old.filter((x) => !items.includes(x)) : old,
557 );
558 });
559 };
560
561 /**
562 * Remove items from an array that match a `filter`.
563 * If the query or path doesn't exist, the updater is skipped.
564 */
565 arrayRemove = <Data>(
566 queryKey: QueryKeyAndFn<Data[] | null>,
567 filter: (
568 item: Data,
569 index: number,
570 ) => boolean,
571 ) => {
572 const prev = this.#get(queryKey);
573 if (!prev || !Array.isArray(prev)) return;
574
575 const newArray = prev.filter((item, index) => !filter(item, index));
576 this.#set(queryKey, newArray);
577 this.#onRestore(() => {
578 // TODO: splice items back in case original changed
579 this.#set(queryKey, prev);
580 });
581 };
582
583 /**
584 * Filter items to just include items that match `filter`. This is the inverse of `arrayRemove`
585 * If the query or path doesn't exist, the updater is skipped.
586 */
587 arrayFilter = <Data>(
588 queryKey: QueryKeyAndFn<Data[] | null>,
589 filter: (
590 item: Data,
591 index: number,
592 ) => boolean,
593 ) => {
594 const prev = this.#get(queryKey);
595 if (!prev || !Array.isArray(prev)) return;
596
597 const newArray = prev.filter(filter);
598 this.#set(queryKey, newArray);
599 this.#onRestore(() => {
600 // TODO: splice items back in case original changed
601 this.#set(queryKey, prev);
602 });
603 };
604
605 /**
606 * Update items in an array that match a `filter`.
607 * If the query or path doesn't exist, the updater is skipped.
608 */
609 arrayUpdate = <Data>(
610 queryKey: QueryKeyAndFn<Data[] | null>,
611 {
612 filter,
613 update,
614 }: {
615 filter?: (item: Data, index: number) => boolean;
616 update: (item: Data) => Data;
617 },
618 ): { inserted: boolean } => {
619 return this.arrayUpsert(queryKey, { filter: filter ?? (() => true), update });
620 };
621
622 /**
623 * Update items in an array that match a `filter`, or insert a new one if there was no match.
624 * If the query doesn't exist, the updater is skipped.
625 */
626 arrayUpsert = <Data>(
627 queryKey: QueryKeyAndFn<Data[] | null>,
628 {
629 filter,
630 update,
631 insert,
632 }:
633 & {
634 filter: (item: Data, index: number) => boolean;
635 }
636 & RequireAtLeastOneKey<{
637 /** Defaults to the identity function */
638 update: (item: Data) => Data;
639 /** Defaults to not inserting */
640 insert: () => Data;
641 }>,
642 ): { inserted: boolean } => {
643 const prev = this.#get(queryKey);
644 if (!prev || !Array.isArray(prev)) return { inserted: false };
645
646 let matched = false;
647 const newArray = prev.map((item, index) => {
648 if (filter(item, index)) {
649 matched = true;
650 return update ? update(item) : item;
651 } else {
652 return item;
653 }
654 });
655 let inserted = false;
656 if (!matched && insert) {
657 newArray.push(insert());
658 inserted = true;
659 }
660 this.#set(queryKey, newArray);
661 this.#onRestore(() => {
662 // TODO: splice items back in case original changed
663 this.#set(queryKey, prev);
664 });
665 return { inserted };
666 };
667
668 /**
669 * Insert item(s) at a specific index in an array.
670 * If the query or path doesn't exist, the updater is skipped.
671 */
672 arrayInsertIndex = <Data>(
673 queryKey: QueryKeyAndFn<Data[] | null>,
674 index: number,
675 ...items: Data[]
676 ) => {
677 const prev = this.#get(queryKey);
678 if (!prev || !Array.isArray(prev)) return;
679
680 const newArray = [
681 ...prev.slice(0, index),
682 ...items,
683 ...prev.slice(index),
684 ];
685 this.#set(queryKey, newArray);
686 this.#onRestore(() => {
687 // TODO: splice items back in case original changed
688 this.#set(queryKey, prev);
689 });
690 };
691
692 /**
693 * Remove a query from the cache entirely.
694 */
695 removeQuery = (queryKey: QueryKeyAndFn) => {
696 const prev = this.#get(queryKey);
697 if (!prev) return;
698
699 this.#client.removeQueries({ queryKey: queryKey.queryKey, exact: true });
700 this.#onRestore(() => {
701 this.#set(queryKey, prev);
702 });
703 };
704}
705
706export function queryClientOptimisticHelpers(
707 client: QueryClient,
708): (e: OptimisticEvents) => TanstackQueryOptimisticHelpers {
709 return (events) => new TanstackQueryOptimisticHelpers(client, events);
710}