The personal website of @erikwittern

Selectively Listen to Firestore Queries with Zustand

September 2nd 2022

We persist data in Firestore at Coup. One nice feature of Firestore is listening to real-time updates of results of a query. In our React application, this enables a data-flow pattern where we listen to some Firestore documents and as soon as we commit updates to them, our UI reflects the latest state (which is in sync with the DB).

However, there is one challenge: we often want to be able to either listen to changes across all queried documents, or only listen to changes in one specific document in different places in our application. In the latter case, we want our UI to re-render only when that one document is changed, but not if any other document changes. We want to achieve this without having to setup listeners for individual Firestore documents, though - in theory, one listener across all documents should be enough!

This can be achieved using the Zustand state management library, relying on its ability to have multiple, isolated stores. We lay out our approach here in TypeScript. Let's assume we are interested in updates to Project documents. First, we define a store for keeping track of all projects. It allows us, for example, to render an up-to-date list of projects:

1
// little type helper for defining stores:
2
type Store<T> = { data: null | T, setData: (data: null | T) => void };
3


4
// store keeping track of all projects:
5
const useProjectsStore = create<Store<Project[]>>(set => ({
6
  data: null,
7
  setData: (data) =>  set(state => ({...state, data}))
8
}));

Next, and this is where things get interesting, we create a Map of stores for each individual project, and a helper function that for a given id of a project returns its corresponding store, or creates a new one if needed:

1
// type `UseBoundStore<StoreAPI<...>>` stems from Zustand:
2
const projectStores = new Map<string, UseBoundStore<StoreApi<Store<Project>>>>()
3


4
/**
5
 * Returns the ProjectStore for the project with the given `projectId`. If that
6
 * store does not yet exist, initializes it first.
7
 */
8
function getProjectStore(projectId: string) {
9
  if (!projectStores.has(projectId)) {
10
    projectStores.set(
11
      projectId,
12
      create<Store<Project>>((set) => ({
13
        data: null,
14
        setData: (data) => set((state) => ({ ...state, data })),
15
      }))
16
    );
17
  }
18
  return projectStores.get(projectId)
19
}

Now we have a setup where one Zustand store keeps track of a whole list of projects, and where multiple additional stores keep track the state of individual projects. To actually put data into these stores with a Firestore query, we create a ProjectsQuery component. It doesn't render anything but is solely responsible for setting up a real-time query and updating stores as data changes. We position ProjectsQuery at the root of our Rect application.

1
// ...store setup from above...
2


3
function ProjectsQuery () {
4
  const setProjects = useProjectsStore(state => state.data);
5


6
  // useEffect with an empty dependency array for effects on mount and unmount:
7
  useEffect(() => {
8
    const q = query(collection(db, 'projects'), where(/* some constraints */));
9
    const unsubscribe = onSnapshot(q, (snapshot) => {
10
      // first, set latest state of all projects in projectsStore:
11
      const projects: Project[] = [];
12
      snapshot.forEach(doc => projects.push(doc.data() as Project));
13
      setProjects(projects);
14


15
      // second, selectively update individual project stores:
16
      snapshot.docChanges().forEach((change) => {
17
        const project = change.doc.data() as Project;
18
        const projectStore = getProjectStore(project.id);
19
        if (change.type === 'added' || change.type === 'modified') {
20
          projectStore.setState((state) => ({ ...state, project }));
21
        } else if (change.type === 'removed') {
22
          projectStore.setState((state) => ({ ...state, project: null }));
23
        }
24
      });
25
    });
26


27
    // cleanup & unsubscribe the real-time query when this component unmounts:
28
    return () => {
29
      setProjects(null);
30
      projectStores.forEach((store) => store.destroy());
31
      projectStores.clear();
32
      unsubscribe()
33
    };
34
  }, []);
35


36
  return null;
37
}

Now, to consume data from these stores, we create two hooks, which abstract away the usage of Zustand:

1
// Hook for consuming a live-updating list of all queried-for projects:
2
export const useProjects() => useProjectsStore((state) => state.data);
3


4
// Hook for consuming a single live-updating project:
5
export const useProject = (projectId: string) =>
6
  getProjectStore(projectId)((state) => state.project);

One great property of the useProject hook is that it will cause a component to re-render if the project with the given projectId changes, but will not do so if any of the other projects change. A typical usage would be:

1
function ExampleComponent({ projectId }) {
2
  const project = useProject(projectId);
3


4
  async function handleNameUpdate() {
5
    await updateDoc(doc(db, 'projects', projectId), {
6
      name: 'new name'
7
    });
8
  }
9


10
  return (
11
    <div>
12
      <button type="button" onClick={handleNameUpdate}>Update name</button>
13
      <p>{project.name}<p>
14
    </div>
15
  )
16
}

Here, we use Firestore's updateDoc function to update a single project document. By using our useProject hook, the component will re-render once the change is performed. If ExampleComponent was used multiple times with different projectIds, they would not re-render as a result.