relay_hooks_with_refetch_status.ts

import { useCallback, useEffect, useRef, useState } from "react";
import {
  usePaginationFragment,
  useRefetchableFragment,
  useRelayEnvironment,
} from "react-relay";
import type { KeyType } from "react-relay/relay-hooks/helpers";
import type {
  DataID,
  Environment,
  FetchQueryFetchPolicy,
  GraphQLTaggedNode,
  MutationConfig,
  MutationParameters,
  OperationType,
  RequestDescriptor,
  Variables,
  VariablesOf,
} from "relay-runtime";
import {
  commitMutation,
  ConnectionHandler,
  createOperationDescriptor,
  createReaderSelector,
  fetchQuery,
  getFragment,
  getRefetchMetadata,
  getSelector,
} from "relay-runtime";
import { RelayEnvironment } from "../RelayEnvironment";

// TODO: Remove once https://github.com/facebook/relay/issues/4526 is resolved
export function useRefetchableFragmentWithRefetchStatus<
  TQuery extends OperationType,
  TKey extends KeyType,
>(fragmentInput: GraphQLTaggedNode, parentFragmentRef: TKey) {
  const [data, relayRefetch] = useRefetchableFragment(
    fragmentInput,
    parentFragmentRef,
  );

  const environment = useRelayEnvironment();
  const [isRefetching, setIsRefetching] = useState(false);

  const currentQueryRef = useRef(0);
  const queryCounterRef = useRef(0);

  const dataRef = useRef(data);
  dataRef.current = data;
  const relayRefetchRef = useRef(relayRefetch);
  relayRefetchRef.current = relayRefetch;
  const fragmentInputRef = useRef(fragmentInput);
  fragmentInputRef.current = fragmentInput;

  const isMountedRef = useRef(true);
  useEffect(() => {
    isMountedRef.current = true;
    return () => {
      isMountedRef.current = false;
    };
  }, []);

  const refetch = useCallback(
    (
      variables: Partial<VariablesOf<TQuery>> = {},
      {
        fetchPolicy,
        silent = false,
      }: { fetchPolicy: FetchQueryFetchPolicy; silent?: boolean } = {
        fetchPolicy: "store-or-network",
      },
    ) => {
      if (!silent) {
        setIsRefetching(true);
      }

      const data = dataRef.current;
      const relayRefetch = relayRefetchRef.current;
      const queryId = ++queryCounterRef.current;

      const fragmentNode = getFragment(fragmentInputRef.current);
      const fragmentSelector = getSelector(fragmentNode, parentFragmentRef);

      let parentVariables: Variables;
      let fragmentVariables: Variables;
      if (fragmentSelector == null) {
        parentVariables = {};
        fragmentVariables = {};
      } else if (fragmentSelector.kind === "PluralReaderSelector") {
        parentVariables = fragmentSelector.selectors[0]?.owner.variables ?? {};
        fragmentVariables = fragmentSelector.selectors[0]?.variables ?? {};
      } else {
        parentVariables = fragmentSelector.owner.variables;
        fragmentVariables = fragmentSelector.variables;
      }

      const { refetchableRequest, refetchMetadata } = getRefetchMetadata(
        fragmentNode,
        "",
      );

      const refetchVariables = {
        ...parentVariables,
        ...fragmentVariables,
        ...variables,
      };

      const { identifierInfo } = refetchMetadata;
      if (
        identifierInfo != null &&
        !Object.hasOwn(variables, identifierInfo.identifierQueryVariableName)
      ) {
        const identifierValue =
          identifierInfo?.identifierField != null &&
          data != null &&
          typeof data === "object"
            ? (data as unknown as Record<string, unknown>)[
                identifierInfo.identifierField
              ]
            : null;

        if (typeof identifierValue !== "string") {
          console.warn(
            `Expected result to have a string '${identifierInfo.identifierField}' in order to refetch, got '${identifierValue}'.`,
          );
        }

        Object.assign(refetchVariables, {
          [identifierInfo.identifierQueryVariableName]: identifierValue,
        });
      }

      const refetchQuery = createOperationDescriptor(
        refetchableRequest,
        refetchVariables,
        { force: true },
      );

      fetchQuery(
        environment,
        refetchableRequest,
        refetchQuery.request.variables,
        { fetchPolicy },
      ).subscribe({
        complete: () => {
          if (currentQueryRef.current > queryId || !isMountedRef.current) {
            return;
          }

          currentQueryRef.current = queryId;
          relayRefetch(refetchQuery.request.variables, {
            fetchPolicy: "store-only",
          });

          if (queryCounterRef.current > queryId) {
            return;
          }

          setIsRefetching(false);
        },
        error: () => {
          if (queryCounterRef.current > queryId) {
            return;
          }

          setIsRefetching(false);
        },
      });
    },
    [environment, parentFragmentRef],
  );

  return [data, refetch, isRefetching] as const;
}

// TODO: Remove once https://github.com/facebook/relay/issues/4526 is resolved
export function usePaginationFragmentWithRefetchStatus<
  TQuery extends OperationType,
  TKey extends KeyType,
>(fragmentInput: GraphQLTaggedNode, parentFragmentRef: TKey) {
  const data = usePaginationFragment(fragmentInput, parentFragmentRef);

  const environment = useRelayEnvironment();
  const [isRefetching, setIsRefetching] = useState(false);

  const currentQueryRef = useRef(0);
  const queryCounterRef = useRef(0);

  const dataRef = useRef(data);
  dataRef.current = data;
  const fragmentInputRef = useRef(fragmentInput);
  fragmentInputRef.current = fragmentInput;

  const isMountedRef = useRef(true);
  useEffect(() => {
    isMountedRef.current = true;
    return () => {
      isMountedRef.current = false;
    };
  }, []);

  const refetch = useCallback(
    (
      variables: Partial<VariablesOf<TQuery>> = {},
      {
        fetchPolicy,
        silent = false,
      }: { fetchPolicy: FetchQueryFetchPolicy; silent?: boolean } = {
        fetchPolicy: "store-or-network",
      },
    ) => {
      if (!silent) {
        setIsRefetching(true);
      }

      const { data, refetch } = dataRef.current;
      const queryId = ++queryCounterRef.current;

      const fragmentNode = getFragment(fragmentInputRef.current);
      const fragmentSelector = getSelector(fragmentNode, parentFragmentRef);

      let parentVariables: Variables;
      let fragmentVariables: Variables;
      if (fragmentSelector == null) {
        parentVariables = {};
        fragmentVariables = {};
      } else if (fragmentSelector.kind === "PluralReaderSelector") {
        parentVariables = fragmentSelector.selectors[0]?.owner.variables ?? {};
        fragmentVariables = fragmentSelector.selectors[0]?.variables ?? {};
      } else {
        parentVariables = fragmentSelector.owner.variables;
        fragmentVariables = fragmentSelector.variables;
      }

      const { refetchableRequest, refetchMetadata } = getRefetchMetadata(
        fragmentNode,
        "",
      );

      const refetchVariables = {
        ...parentVariables,
        ...fragmentVariables,
        ...variables,
      };

      const { identifierInfo } = refetchMetadata;
      if (
        identifierInfo != null &&
        !Object.hasOwn(variables, identifierInfo.identifierQueryVariableName)
      ) {
        const identifierValue =
          identifierInfo?.identifierField != null &&
          data != null &&
          typeof data === "object"
            ? (data as unknown as Record<string, unknown>)[
                identifierInfo.identifierField
              ]
            : null;

        if (typeof identifierValue !== "string") {
          console.warn(
            `Expected result to have a string '${identifierInfo.identifierField}' in order to refetch, got '${identifierValue}'.`,
          );
        }

        Object.assign(refetchVariables, {
          [identifierInfo.identifierQueryVariableName]: identifierValue,
        });
      }

      const refetchQuery = createOperationDescriptor(
        refetchableRequest,
        refetchVariables,
        { force: true },
      );

      fetchQuery(
        environment,
        refetchableRequest,
        refetchQuery.request.variables,
        { fetchPolicy },
      ).subscribe({
        complete: () => {
          if (currentQueryRef.current > queryId || !isMountedRef.current) {
            return;
          }

          currentQueryRef.current = queryId;
          refetch(refetchQuery.request.variables, {
            fetchPolicy: "store-only",
          });

          if (queryCounterRef.current > queryId) {
            return;
          }

          setIsRefetching(false);
        },
        error: () => {
          if (queryCounterRef.current > queryId) {
            return;
          }

          setIsRefetching(false);
        },
      });
    },
    [environment, parentFragmentRef],
  );

  return { ...data, refetch, isRefetching };
}