Custom useInfiniteQuery cache shape

[vjuux]

Good day react query community, I'm looking for a way to customize the root shape of the useInfiniteQuery cache.

Instead of

data: {
   pages: Object[][];
   pageParams: {};
}

to something like

data: {
   items: id[];
   itemsById: {
      [id]: Object
   }
   xnext: string // used for fetching further batch of objects
}

when fetchNextPage is being hit, next batch of objects is being indexified and added to itemsById & items.
I find this shape easier to work with when using WS, currently when receiving an update/ delete/ create it takes iterations over pages to identify the object in order to manipulate, with this new shape it would flatten things out and simplify query cache interaction. Currently it's an extremely convoluted and inefficient queryClient.setQueriesData() manipulation.

Initial idea was to create a custom useInfiniteQuery by creating a customQueryBehavior and customQueryObserver that would be passed in useBaseQuery. Where buildNewPages and finalPromise would massage the payload accordingly. It looks like the useBaseQuery is not exported for a reason?

This is being used for an infinite scroll use case where tracking pages is not necessary.
Is this even a good idea or I'm stretching the RQ too far here?

@TkDodo 🙏🥺

[Arian-Hamdi]

Hey @vjuux,
Did you find any solution for this?

In our case, we have an infinite scroll chat history, and we need to update a specific ID frequently because of streaming. Looping over all the messages is not efficient. We were thinking about creating a map to be able to find the desired message's index with O(1) complexity, but the shape of the cache in useInfiniteQuery doesn't allow it. We would have to create that map index inside each page individually instead of something global.

That would be great if you can take a look at this @TkDodo @tannerlinsley

[allohamora]

#3661 (comment)

[allohamora]

You can try to make a solution like this:

import { RefetchOptions, useQuery, useQueryClient } from '@tanstack/react-query';

export type RealtimeInfinitePage<TItem, TPageParams> = {
  items: TItem[];
  nextPageParams?: TPageParams;
  total?: number;
};

export type RealtimeInfiniteData<TItem, TPageParams> = RealtimeInfinitePage<TItem, TPageParams> & { lastItem?: TItem };

export type RealtimeInfiniteUpdater<TItem, TPageParams> = (
  oldData: RealtimeInfiniteData<TItem, TPageParams> | undefined,
) => RealtimeInfiniteData<TItem, TPageParams> | undefined;

export type UseRealtimeInfiniteQueryOptions<TItem, TPageParams> = {
  queryKey: unknown[];
  fetchPage: (
    pageParams: TPageParams | undefined,
    signal?: AbortSignal,
  ) => Promise<RealtimeInfinitePage<TItem, TPageParams>>;
  mergeItems?: (oldItems: TItem[], newItems: TItem[]) => TItem[];
};

export const concatItems = <TItem>(oldItems: TItem[], newItems: TItem[]): TItem[] => [...oldItems, ...newItems];

export const strictConcatItems =
  <TItem, TKey extends keyof TItem = keyof TItem>(idField: TKey) =>
  (oldItems: TItem[], newItems: TItem[]): TItem[] => {
    const oldIds = new Set(oldItems.map((item) => item[idField]));
    const duplicates = newItems.filter((item) => oldIds.has(item[idField]));

    if (duplicates.length > 0) {
      console.error(
        new Error(
          'useRealtimeInfiniteQuery: mergeItems found id(s) present in both the already-loaded items and the ' +
            'freshly fetched page - pass a mergeItems that resolves the overlap',
        ),
        { idField, duplicates },
      );
    }

    return concatItems(oldItems, newItems);
  };

// purpose-built alternative to useInfiniteQuery for cursor-based infinite scroll that also has to
// absorb realtime pushes into the same cache. useInfiniteQuery's internal pages/pageParams array
// makes patching in one realtime item awkward (rebuilding page boundaries by hand on every event);
// a flat `items` list turns that into a single setQueryData call instead.
//
// deliberately not at parity with useInfiniteQuery: invalidate/refetch always collapses back to
// page 1 (no full-depth refetch), no fetchPreviousPage, no maxPages eviction, no automatic dedup
// at the page boundary (mergeItems defaults to plain concat - same as useInfiniteQuery itself;
// pass strictConcatItems or your own mergeItems if pages can overlap)
export const useRealtimeInfiniteQuery = <TItem, TPageParams = string>({
  queryKey,
  fetchPage,
  mergeItems = concatItems,
}: UseRealtimeInfiniteQueryOptions<TItem, TPageParams>) => {
  const queryClient = useQueryClient();

  // the discriminator goes first, not last (i.e. not [...queryKey, 'next-page']) - a trailing
  // discriminator would still start with queryKey, so filters like invalidateQueries({ queryKey })
  // would prefix-match and silently sweep up this internal query too
  const nextPageQueryKey = ['use-realtime-infinite-query-next-page', ...queryKey];

  const query = useQuery({
    queryKey,
    queryFn: async ({ signal }): Promise<RealtimeInfiniteData<TItem, TPageParams>> => {
      const page = await fetchPage(undefined, signal);

      // fetchPage might resolve instead of throwing even once aborted (e.g. it doesn't forward
      // the signal all the way to the real request) - a superseding refetch already owns clearing
      // next-page state below, so a late-resolving stale call must not fire that side effect again
      // after the superseding call has already re-established valid next-page state
      if (!signal.aborted) {
        // a fresh page 1 means any previous "load more" attempt no longer applies to this
        // pagination session - without this, a stale isNextPageError/nextPageError from before
        // this fetch (a failed fetchNextPage, or one inherited from before a remount) would keep
        // showing even though nothing has failed yet in this session
        queryClient.resetQueries({ queryKey: nextPageQueryKey, exact: true });
      }

      // this return unconditionally replaces the cache, so a realtime setQueryData call that lands
      // while this fetch is in flight would be overwritten - callers must not enable their realtime
      // subscription until this fetch has already settled (see the two useXUpdateSubscription callers)
      return { ...page, lastItem: page.items.at(-1) };
    },
    // default staleTime (0): remounting (e.g. browser back/forward) refetches page 1 and resets
    // pagination depth back to a single page - accepted tradeoff in exchange for always-fresh data
    refetchOnWindowFocus: false,
    refetchOnReconnect: false,
  });

  // a plain useQuery instead of a useMutation: only a real queryFn gets an injected AbortSignal
  // (auto-cancelled on unmount) and the framework's default retry behavior - a mutation gets neither
  const nextPageQuery = useQuery({
    queryKey: nextPageQueryKey,
    queryFn: async ({ signal }) => {
      const current = queryClient.getQueryData<RealtimeInfiniteData<TItem, TPageParams>>(queryKey);
      if (current?.nextPageParams === undefined) return null;

      const page = await fetchPage(current.nextPageParams, signal);

      // fetchPage might resolve instead of throwing even once aborted (e.g. it doesn't forward
      // the signal all the way to the real request) - a superseding call already owns the merge,
      // so applying this one too would double-append the same page
      if (signal.aborted) return page;

      // useQuery has no onSuccess callback in v5, so the merge happens inline, right after the
      // fetch resolves and before this queryFn returns
      queryClient.setQueryData<RealtimeInfiniteData<TItem, TPageParams>>(queryKey, (old) => {
        if (!old) return old;

        const items = mergeItems(old.items, page.items);
        return { ...page, items, lastItem: items.at(-1) };
      });

      return page;
    },
    // only ever runs when fetchNextPage() calls refetch() below - enabled: false already blocks
    // every automatic refetch path (mount, window focus, reconnect), so there's nothing left for
    // refetchOnWindowFocus/refetchOnReconnect to suppress here
    enabled: false,
    // belt-and-suspenders: redundant while enabled stays false, but keeps this query inert even if
    // that ever changes
    refetchOnWindowFocus: false,
    refetchOnReconnect: false,
    // this cache entry holds no state anything outside the hook should ever read - the resetQueries
    // call above already clears a stale error/status on a fresh page 1, but gcTime: 0 is a cheap
    // backstop: on unmount it's dropped immediately rather than surviving for the default 5 minutes
    gcTime: 0,
  });

  const setQueryData = (updater: RealtimeInfiniteUpdater<TItem, TPageParams>) => {
    queryClient.setQueryData<RealtimeInfiniteData<TItem, TPageParams>>(queryKey, updater);
  };

  return {
    items: query.data?.items,
    total: query.data?.total,
    lastItem: query.data?.lastItem,
    hasNextPage: query.data?.nextPageParams !== undefined,

    // page 1 only - never affected by anything happening to nextPageQuery below
    isFirstPagePending: query.isPending, // still loading, no data at all yet
    isFirstPageRefetching: query.isRefetching, // data already exists, but page 1 is being re-fetched (e.g. a remount)
    isFirstPageError: query.isError,
    firstPageError: query.error,
    // discards any pages loaded via fetchNextPage, collapsing back to a single fresh page 1 -
    // only safe to call where there's nothing loaded yet to lose (e.g. an empty/error-state retry).
    // cancelRefetch: false for the same reason as fetchNextPage below - an overlapping call (e.g. a
    // double-clicked retry/refresh once page 1 already has data) dedupes onto the in-flight retryer
    // instead of cancelling and restarting it, so the signal.aborted guard above is never bypassed
    // by a second concurrent execution of this queryFn
    resetToFirstPage: (options?: RefetchOptions) => query.refetch({ cancelRefetch: false, ...options }),

    // "load more" only - never affects page 1's own state above, and never touches `items` on failure
    isFetchingNextPage: nextPageQuery.isFetching,
    isNextPageError: nextPageQuery.isError,
    nextPageError: nextPageQuery.error,
    // cancelRefetch: false - a call while one's already in flight dedupes onto the same request
    // instead of cancelling and restarting it, so the merge above never runs twice concurrently.
    // only matters from the second "load more" onward: query-core's fetch() only takes the
    // cancel-and-restart path when the query already has data, so on the very first call (before
    // nextPageQuery has ever resolved) it dedupes onto the in-flight retryer either way
    fetchNextPage: (options?: RefetchOptions) => nextPageQuery.refetch({ cancelRefetch: false, ...options }),

    setQueryData,
  };
};

type Item = { id: string; sortKey: string };
type Event = { type: 'create' | 'update'; item: Item } | { type: 'delete'; id: string };
type Data = RealtimeInfiniteData<Item, string>;

// mirrors the server's order-by: sortKey desc
const compareItems = (a: Item, b: Item) => b.sortKey.localeCompare(a.sortKey);

const applyEvent = (event: Event, oldData?: Data): Data | undefined => {
  if (!oldData) return oldData; // let the initial fetch populate the cache

  const itemById = new Map(oldData.items.map((item) => [item.id, item]));

  // lastItem is the page cursor, not derived from `items` - it only advances on a real fetch
  // (initial load or fetchNextPage), so it marks how far we've actually loaded. We use it here to
  // decide whether an event we don't already have cached falls inside that loaded range (apply it)
  // or beyond it (ignore it - a future page will pick it up on its own)
  const isWithinLoadedRange = (item: Item) => !oldData.lastItem || compareItems(item, oldData.lastItem) <= 0;

  let totalDelta = 0;

  switch (event.type) {
    case 'create':
    case 'update': {
      const isAlreadyCached = itemById.has(event.item.id);

      // out-of-range only disqualifies an item we don't already know about - an already-cached
      // item's update is always safe to apply regardless of where it now sorts
      if (isWithinLoadedRange(event.item) || isAlreadyCached) {
        if (!isAlreadyCached) totalDelta = 1;
        itemById.set(event.item.id, event.item);
      }
      break;
    }
    case 'delete':
      if (itemById.delete(event.id)) totalDelta = -1;
      break;
  }

  const items = Array.from(itemById.values()).sort(compareItems);

  // lastItem is deliberately left as-is here, not recomputed from items.at(-1): a realtime
  // delete/create can change the tail of `items` without the cursor itself having moved, so
  // rederiving it from the post-merge list would wrongly shrink or grow the loaded range
  return { ...oldData, items, total: (oldData.total ?? 0) + totalDelta };
};

const useItemsQuery = () => {
  const query = useRealtimeInfiniteQuery({
    queryKey: ['items'],
    fetchPage: (cursor, signal) => fetchItemsPage(cursor, signal),
    mergeItems: strictConcatItems('id'),
  });

  // don't subscribe while a fetch (initial load or a retry) is in flight - its unconditional cache
  // replace would stomp any event that lands in the meantime
  const isFirstPageSettled = !query.isFirstPagePending && !query.isFirstPageRefetching && !query.isFirstPageError;

  useSubscription({
    onData: (event) => query.setQueryData((oldData) => applyEvent(event, oldData)),
    enabled: isFirstPageSettled,
  });

  return query;
};