createPreloadedQuery.ts

import { useCallback, useEffect, useMemo } from 'react';
// eslint-disable-next-line no-restricted-imports
import { loadQuery as relayLoadQuery, LoadQueryOptions } from 'react-relay';
import {
  GraphQLTaggedNode,
  OperationType,
  usePreloadedQuery as usePreloadedQueryDataLayer,
  useRelayEnvironment,
  VariablesOf,
} from '@cbhq/data-layer';

function useCleanUpOnUnmount(cb: () => void) {
  // 2022-12-15 andrei-calazans: callback is not intended to be stable
  // and we truly only want to clean up on parent component unmount.
  // eslint-disable-next-line
  return useEffect(() => () => cb(), []);
}

/*
 * What's createPreloadedQuery?
 *
 * The goal of this factory function is to make it easy
 * to create a preloaded query and share its query reference.
 *
 * Without this you would need to store the query reference in a
 * React store somewhere like Context without needing React's
 * reactivity.
 *
 * This API stores the shared query reference in the factory
 * function scope and takes care of cleaning it up when the consumer
 * usePreloadedQuery gets unmounted.
 *
 * It also takes care of calling Relay's loadQuery to retrieve
 * the query reference when a prefetch did not occur - this is
 * typical in scenarios where you navigate to a screen
 * through deeplinking without prefetching via loadQuery.
 *
 * */
export function createPreloadedQuery<TQuery extends OperationType>(
  query: GraphQLTaggedNode,
) {
  let preloadedQueryRef: ReturnType<typeof relayLoadQuery> | null = null;

  const useLoadQuery = () => {
    const relayEnvironment = useRelayEnvironment();
    const loadQuery = useCallback(
      (variables: VariablesOf<TQuery>, options?: LoadQueryOptions) => {
        preloadedQueryRef = relayLoadQuery(
          relayEnvironment,
          query,
          variables,
          options,
        );
        return preloadedQueryRef;
      },
      [relayEnvironment],
    );

    return loadQuery;
  };

  const usePreloadedQuery = (
    variables: VariablesOf<TQuery>,
    options?: LoadQueryOptions,
  ): TQuery['response'] => {
    useCleanUpOnUnmount(() => {
      if (preloadedQueryRef && 'dispose' in preloadedQueryRef) {
        preloadedQueryRef.dispose();
        preloadedQueryRef = null;
      }
    });

    const internalLoadQuery = useLoadQuery();
    const queryRef = useMemo(
      () => preloadedQueryRef ?? internalLoadQuery(variables, options),
      // 2022-12-15 andrei-calazans: we don't want to include
      // variables and options as a depedency since they are often
      // unstable and inlined objects.
      // eslint-disable-next-line react-hooks/exhaustive-deps
      [preloadedQueryRef],
    );

    return usePreloadedQueryDataLayer(query, queryRef);
  };

  return {
    usePreloadedQuery,
    useLoadQuery,
  };
}