Apollo query hook that supports React.Suspense

useSuspenseQuery.ts

import { ApolloQueryResult, QueryOptions, QueryResult, TypedDocumentNode, useApolloClient, useQuery } from "@apollo/client";
import uuid from "uuid/v5";

const uuidNamespace = uuid("suspense.queries.utils", uuid.DNS);
const suspenseQueryStore: Record<string, { read: () => ApolloQueryResult<unknown> }> = {};

type SuspendQueryResult<Q, V> = Pick<QueryResult<Q, V>, "networkStatus" | "variables"> & { data: Q };

/** Works with `TypedDocumentNode` which can be generated using `typed-document-node` codegen plugin. */
export function useSuspenseQuery<Q, V>(
  query: TypedDocumentNode<Q, V>,
  options: Omit<QueryOptions<V, Q>, "query"> = {}
): SuspendQueryResult<Q, V> {
  const apolloClient = useApolloClient();

  const uniqueQueryId = uuid(JSON.stringify([query, options.variables]), uuidNamespace);

  if (!suspenseQueryStore[uniqueQueryId]) {
    const mergedOptions: QueryOptions<V, Q> = { ...options, query };
    const promise = apolloClient.query<Q, V>(mergedOptions);
    suspenseQueryStore[uniqueQueryId] = wrapPromise(promise);
  }

  // this line throws on loading (Suspense) or error (ErrorBoundary)
  suspenseQueryStore[uniqueQueryId].read();

  // IMPORTANT: useQuery() must follow read(), otherwise memory leak
  const { data, variables, networkStatus } = useQuery(query, options);
  if (!data) {
    throw Error("We don't expect empty data at this point");
  }

  return {
    data,
    networkStatus,
    variables,
  };
}

function wrapPromise<T extends unknown>(promise: Promise<T>) {
  let status: "pending" | "success" | "error" = "pending";
  let resultOrError: T;

  const suspender = promise.then(
    (r) => {
      status = "success";
      resultOrError = r;
    },
    (e) => {
      status = "error";
      resultOrError = e;
    }
  );

  const read = () => {
    switch (status) {
      case "pending":
        throw suspender;
      case "error":
        throw resultOrError;
      default:
        return resultOrError;
    }
  };

  return { read };
}