ra-core/src/dataProvider/useGetManyAggregate.ts

Core components of react-admin, a frontend Framework for building admin applications on top of REST services, using ES6, React

import { useEffect, useMemo, useRef } from 'react';
import {
    notifyManager,
    QueryClient,
    QueryKey,
    useQueryClient,
    useQuery,
    UseQueryOptions,
} from '@tanstack/react-query';
import union from 'lodash/union.js';

import { UseGetManyHookValue } from './useGetMany';
import { Identifier, RaRecord, GetManyParams, DataProvider } from '../types';
import { useDataProvider } from './useDataProvider';
import { useEvent } from '../util';

/**
 * Call the dataProvider.getMany() method and return the resolved result
 * as well as the loading state.
 *
 * The return value updates according to the request state:
 *
 * - start: { isPending: true, isFetching: true, refetch }
 * - success: { data: [data from response], isPending: false, isFetching: false, refetch }
 * - error: { error: [error from response], isPending: false, isFetching: false, refetch }
 *
 * This hook will return the cached result when called a second time
 * with the same parameters, until the response arrives.
 *
 * This hook aggregates and deduplicates calls to the same resource, so for instance, if an app calls:
 *
 * useGetManyAggregate('tags', [1, 2, 3]);
 * useGetManyAggregate('tags', [3, 4]);
 *
 * during the same tick, the hook will only call the dataProvider once with the following parameters:
 *
 * dataProvider.getMany('tags', [1, 2, 3, 4])
 *
 * @param resource The resource name, e.g. 'posts'
 * @param {Params} params The getMany parameters { ids, meta }
 * @param {Object} options Options object to pass to the dataProvider.
 * @param {boolean} options.enabled Flag to conditionally run the query. If it's false, the query will not run
 * @param {Function} options.onSuccess Side effect function to be executed upon success, e.g. { onSuccess: { refresh: true } }
 * @param {Function} options.onError Side effect function to be executed upon failure, e.g. { onError: error => notify(error.message) }
 *
 * @typedef Params
 * @prop params.ids The ids to get, e.g. [123, 456, 789]
 * @prop params.meta Optional meta parameters

 * @returns The current request state. Destructure as { data, error, isPending, isFetching, refetch }.
 *
 * @example
 *
 * import { useGetManyAggregate, useRecordContext } from 'react-admin';
 *
 * const PostTags = () => {
 *     const record = useRecordContext();
 *     const { data, isPending, error } = useGetManyAggregate('tags', { ids: record.tagIds });
 *     if (isPending) { return <Loading />; }
 *     if (error) { return <p>ERROR</p>; }
 *     return (
 *          <ul>
 *              {data.map(tag => (
 *                  <li key={tag.id}>{tag.name}</li>
 *              ))}
 *          </ul>
 *      );
 * };
 */
export const useGetManyAggregate = <
    RecordType extends RaRecord = any,
    ErrorType = Error,
>(
    resource: string,
    params: Partial<GetManyParams<RecordType>>,
    options: UseGetManyAggregateOptions<RecordType, ErrorType> = {}
): UseGetManyHookValue<RecordType, ErrorType> => {
    const dataProvider = useDataProvider();
    const queryClient = useQueryClient();
    const {
        onError = noop,
        onSuccess = noop,
        onSettled = noop,
        enabled,
        ...queryOptions
    } = options;
    const onSuccessEvent = useEvent(onSuccess);
    const onErrorEvent = useEvent(onError);
    const onSettledEvent = useEvent(onSettled);

    const { ids, meta } = params;
    const placeholderData = useMemo(() => {
        const records = (Array.isArray(ids) ? ids : [ids]).map(id =>
            queryClient.getQueryData<RecordType>([
                resource,
                'getOne',
                { id: String(id), meta },
            ])
        );
        if (records.some(record => record === undefined)) {
            return undefined;
        } else {
            return records as RecordType[];
        }
    }, [ids, queryClient, resource, meta]);

    const result = useQuery<RecordType[], ErrorType, RecordType[]>({
        queryKey: [
            resource,
            'getMany',
            {
                ids: (Array.isArray(ids) ? ids : [ids]).map(id => String(id)),
                meta,
            },
        ],
        queryFn: queryParams =>
            new Promise((resolve, reject) => {
                if (!ids || ids.length === 0) {
                    // no need to call the dataProvider
                    return resolve([]);
                }

                // debounced / batched fetch
                return callGetManyQueries({
                    resource,
                    ids,
                    meta,
                    queryKey: queryParams.queryKey,
                    resolve,
                    reject,
                    dataProvider,
                    queryClient,
                    signal:
                        dataProvider.supportAbortSignal === true
                            ? queryParams.signal
                            : undefined,
                });
            }),
        placeholderData,
        enabled: enabled ?? ids != null,
        retry: false,
        ...queryOptions,
    });

    const metaValue = useRef(meta);
    const resourceValue = useRef(resource);

    useEffect(() => {
        metaValue.current = meta;
    }, [meta]);

    useEffect(() => {
        resourceValue.current = resource;
    }, [resource]);

    useEffect(() => {
        if (
            result.data === undefined ||
            result.error != null ||
            result.isFetching
        )
            return;

        // optimistically populate the getOne cache
        (result.data ?? []).forEach(record => {
            queryClient.setQueryData(
                [
                    resourceValue.current,
                    'getOne',
                    { id: String(record.id), meta: metaValue.current },
                ],
                oldRecord => oldRecord ?? record
            );
        });

        onSuccessEvent(result.data);
    }, [
        queryClient,
        onSuccessEvent,
        result.data,
        result.error,
        result.isFetching,
    ]);

    useEffect(() => {
        if (result.error == null || result.isFetching) return;
        onErrorEvent(result.error);
    }, [onErrorEvent, result.error, result.isFetching]);

    useEffect(() => {
        if (result.status === 'pending' || result.isFetching) return;
        onSettledEvent(result.data, result.error);
    }, [
        onSettledEvent,
        result.data,
        result.error,
        result.status,
        result.isFetching,
    ]);

    return result;
};

/**
 * Batch all calls to a function into one single call with the arguments of all the calls.
 *
 * @example
 * let sum = 0;
 * const add = (args) => { sum = args.reduce((arg, total) => total + arg, 0); };
 * const addBatched = batch(add);
 * addBatched(2);
 * addBatched(8);
 * // add will be called once with arguments [2, 8]
 * // and sum will be equal to 10
 */
const batch = fn => {
    let capturedArgs: any[] = [];
    let timeout: ReturnType<typeof setTimeout> | null = null;
    return (arg: any) => {
        capturedArgs.push(arg);
        if (timeout) clearTimeout(timeout);
        timeout = setTimeout(() => {
            timeout = null;
            fn([...capturedArgs]);
            capturedArgs = [];
        }, 0);
    };
};

interface GetManyCallArgs {
    resource: string;
    ids: Identifier[];
    meta?: any;
    queryKey: QueryKey;
    resolve: (data: any[]) => void;
    reject: (error?: any) => void;
    dataProvider: DataProvider;
    queryClient: QueryClient;
    signal?: AbortSignal;
}

/**
 * Resolve all the pending calls with the data they requested, in a single React update.
 *
 * Each pending call belongs to a distinct useQuery. Resolving a promise only makes
 * its query transition two microtasks later, long after any synchronous
 * notifyManager transaction has closed, so resolving the calls one by one makes
 * react-query flush one observer notification per call, and React commits each of
 * them separately. Any child updating its state during the commit phase (e.g. an
 * avatar reporting its image loading status from a layout effect) then turns those
 * commits into nested updates instead of batched ones, and React throws
 * "Maximum update depth exceeded" past 50 of them.
 *
 * Writing the data of every pending call inside a single notifyManager transaction
 * makes react-query flush all the observer notifications at once, so all the
 * consumers get their data in one commit. The calls are then resolved to let their
 * query leave the fetching state.
 */
const resolveCallsWithData = (
    calls: GetManyCallArgs[],
    queryClient: QueryClient,
    getDataForCall: (call: GetManyCallArgs) => any[]
) => {
    const resolutions = calls.map(call => ({
        call,
        data: getDataForCall(call),
    }));
    notifyManager.batch(() => {
        resolutions.forEach(({ call, data }) => {
            const query = queryClient
                .getQueryCache()
                .find({ queryKey: call.queryKey, exact: true });
            // don't resurrect a query that was canceled or removed in the meantime
            if (query?.state.fetchStatus !== 'fetching') return;
            queryClient.setQueryData(call.queryKey, data);
        });
    });
    resolutions.forEach(({ call, data }) => call.resolve(data));
};

const filterDataForCall = (data: any[]) => (call: GetManyCallArgs) =>
    data.filter(record =>
        call.ids.map(id => String(id)).includes(String(record.id))
    );

/**
 * Group and execute all calls to the dataProvider.getMany() method for the current tick
 *
 * Thanks to batch(), this function executes at most once per tick,
 * whatever the number of calls to useGetManyAggregate().
 */
const callGetManyQueries = batch((calls: GetManyCallArgs[]) => {
    const dataProvider = calls[0].dataProvider;
    const queryClient = calls[0].queryClient;

    /**
     * Aggregate calls by resource and meta
     *
     * callsByResourceAndMeta will look like:
     * {
     *     'posts|{"test":true}': [{ resource, ids, resolve, reject, dataProvider, queryClient }, ...],
     *     'posts|{"test":false}': [{ resource, ids, resolve, reject, dataProvider, queryClient }, ...],
     *     tags: [{ resource, ids, resolve, reject, dataProvider, queryClient }, ...],
     * }
     */
    const callsByResourceAndMeta = calls.reduce(
        (acc, callArgs) => {
            const key = `${callArgs.resource}|${JSON.stringify(callArgs.meta)}`;
            if (!acc[key]) {
                acc[key] = [];
            }
            acc[key].push(callArgs);
            return acc;
        },
        {} as { [resource: string]: GetManyCallArgs[] }
    );

    /**
     * For each resource/meta association, aggregate ids and call dataProvider.getMany() once
     */
    Object.keys(callsByResourceAndMeta).forEach(resource => {
        const callsForResource = callsByResourceAndMeta[resource];

        const uniqueResource = callsForResource.reduce(
            (acc, { resource }) => resource || acc,
            '' as string // Should never happen as we always have a resource in callArgs but makes TS happy
        );
        /**
         * Extract ids from queries, aggregate and deduplicate them
         *
         * @example from [[1, 2], [2, null, 3], [4, null]] to [1, 2, 3, 4]
         */
        const aggregatedIds = callsForResource
            .reduce((acc, { ids }) => union(acc, ids), []) // concat + unique
            .filter(v => v != null && v !== ''); // remove null values

        const uniqueMeta = callsForResource.reduce(
            (acc, { meta }) => meta || acc,
            undefined
        );

        if (aggregatedIds.length === 0) {
            // no need to call the data provider if all the ids are null
            resolveCallsWithData(callsForResource, queryClient, () => []);
            return;
        }

        const callThatHasAllAggregatedIds = callsForResource.find(
            ({ ids, signal }) =>
                JSON.stringify(ids) === JSON.stringify(aggregatedIds) &&
                !signal?.aborted
        );
        if (callThatHasAllAggregatedIds) {
            // There is only one call (no aggregation), or one of the calls has the same ids as the sum of all calls.
            // Either way, we can't trigger a new fetchQuery with the same signature, as it's already pending.
            // Therefore, we reply with the dataProvider
            const { dataProvider, resource, ids, meta, signal } =
                callThatHasAllAggregatedIds;

            dataProvider
                .getMany<any>(resource, { ids, meta, signal })
                .then(({ data }) => data)
                .then(
                    data => {
                        // We must then resolve all the pending calls with the data they requested
                        resolveCallsWithData(
                            callsForResource,
                            queryClient,
                            filterDataForCall(data)
                        );
                    },
                    error => {
                        // All pending calls must also receive the error
                        callsForResource.forEach(({ reject }) => {
                            reject(error);
                        });
                    }
                );
            return;
        }

        /**
         * Call dataProvider.getMany() with the aggregatedIds,
         * and resolve each of the promises using the results
         */
        queryClient
            .fetchQuery<any[], Error, any[]>({
                queryKey: [
                    uniqueResource,
                    'getMany',
                    {
                        ids: aggregatedIds.map(id => String(id)),
                        meta: uniqueMeta,
                    },
                ],
                queryFn: queryParams =>
                    dataProvider
                        .getMany<any>(uniqueResource, {
                            ids: aggregatedIds,
                            meta: uniqueMeta,
                            signal:
                                dataProvider.supportAbortSignal === true
                                    ? queryParams.signal
                                    : undefined,
                        })
                        .then(({ data }) => data),
            })
            .then(data => {
                resolveCallsWithData(
                    callsForResource,
                    queryClient,
                    filterDataForCall(data)
                );
            })
            .catch(error =>
                callsForResource.forEach(({ reject }) => reject(error))
            );
    });
});

const noop = () => undefined;

export type UseGetManyAggregateOptions<
    RecordType extends RaRecord = any,
    ErrorType = Error,
> = Omit<UseQueryOptions<RecordType[], ErrorType>, 'queryKey' | 'queryFn'> & {
    onSuccess?: (data: RecordType[]) => void;
    onError?: (error: ErrorType) => void;
    onSettled?: (data?: RecordType[], error?: ErrorType | null) => void;
};