React Query with Firestore subscriptions

[gragland]

Hey folks, I whipped up an example of React Query with Firestore subscriptions that I wanted to share in case it's useful for others. Would love to hear any thoughts on how this could be improved. This would likely be part of a React boilerplate I offer, so trying to make sure it's bug free and easy to understand.

import { getFirestore, onSnapshot, doc, collection, query, where, orderBy } from "firebase/firestore";
import { useQuery, QueryClient, hashQueryKey } from "react-query";
import { firebaseApp } from "./firebase";
const firestore = getFirestore(firebaseApp);

const queryClient = new QueryClient({
  defaultOptions: {
    queries: {
      // No need to refetch because subscriptions keep things up to date
      refetchOnWindowFocus: false,
    },
  },
});

/*** USAGE ***/

// Subscribe to item data
export function useItem(id) {
  return useQuery(
    ["item", id],
    createQuery(() => doc(firestore, "items", id)),
    { enabled: !!id }
  );
}

// Subscribe to all items by owner
export function useItemsByOwner(owner) {
  return useQuery(
    ["items", { owner }],
    createQuery(() =>
      query(
        collection(firestore, "items"),
        where("owner", "==", owner),
        orderBy("createdAt", "desc")
      )
    ),
    { enabled: !!owner }
  );
}

/**** HELPERS ****/

// Create a query function for `useQuery` that subscribes to a Firestore doc or collection
function createQuery(getRef) {
  return async (context) => {
    let firstRun = true;
    let unsubscribe;
    // Wrap `onSnapshot` with a promise because `useQuery` expects a promise
    const data = await new Promise((resolve, reject) => {
      unsubscribe = onSnapshot(
        getRef(),
        // Success callback resolves the promise on the first run
        // For subsequent runs it manually updates the React Query cache
        (response) => {
          const data = format(response);
          if (firstRun) {
            resolve(data);
            firstRun = false;
          } else {
            queryClient.setQueryData(context.queryKey, data);
          }
        },
        // Error callback rejects the promise on the first run
        // We can't manually trigger an error in React Query, so on a subsequent runs we
        // invalidate the query (causes a re-fetch and will reject if error persists)
        (error) => {
          if (firstRun) {
            reject(error);
            firstRun = false;
          } else {
            queryClient.invalidateQueries(context.queryKey);
          }
        }
      );
    });

    // Get query object from React Query
    const query = queryClient.queryCache.find(context.queryKey);
    // Remove existing subscription (if there is one) and store new one on `query` object
    query.__unsubscribe && query.__unsubscribe();
    query.__unsubscribe = unsubscribe;
    // Return data for `useQuery` hook
    return data;
  };
}

// Auto-remove subscriptions when all observing components have unmounted
queryClient.queryCache.subscribe(({ type, query }) => {
  if (type === "observerRemoved" && query.getObserversCount() === 0 && query.__unsubscribe) {
    query.__unsubscribe();
    delete query.__unsubscribe;
  }
});

// Format Firestore response
function format(response) {
  if (response.docs) {
    return response.docs.map(formatDoc);
  } else {
    return formatDoc(response);
  }
}

// Format document (merges `doc.data()` and `doc.id`)
function formatDoc(doc) {
  return doc.exists() === true ? { id: doc.id, ...doc.data() } : null;
}

[0x80]

@gragland I'll soon start working on a new app that combines TanStack Query with Firestore. Finding a good pattern for integrating Firestore subscriptions with state management always seems to be a challenge.

Since this post is fairly old, I wonder; Have you been happy with your solution? Did you have some insights since?