Using Websockets with React Query

I am a massive fan of React Query. It has grown on me to the point that it's now one of the first tools I set up when I start on a new project. If you're not familiar with React Query, at a high level it's a library that helps you manage your data fetching and client-side caching in React applications.

React Query makes it easy to write declarative data fetching code that gets access to all sorts of nice features like request states (isLoading, isFetched, etc), query deduping, query cancellation, cache control and persistence, and a whole lot more — and it gives you the power to manipulate your client-side cache in pretty much any way you might need.

If you just need to fetch some data you can use useQuery with one or two lines of code and be on your way, but as your application grows in complexity you may find yourself wanting to connect to a websocket to receive real-time updates from your server. Fortunately, React Query makes this quite painless.

Querying data

When you're ready to start updating your client via websockets, if you're already using React Query, the good news is that nothing really changes with the way you fetch your data.

const { data } = useQuery({
  queryKey: ['replies', postId],
  queryFn: () => fetchReplies(postId),
});

Here we've got a fetchReplies function that returns a promise that resolves to an array of replies for a given post. We're using the queryKey to identify this query, and we're using the queryFn to call our fetchReplies function.

Subscribing to a websocket

In this example, we're going to use Pusher. I like Pusher because they have an open specification and there are other, self-hostable alternatives that implement the Pusher spec. Meaning you can build on Pusher and if your application scales to the incredibly expensive tiers of Pusher.com, you don't have to rewrite your entire socket architecture to cut some costs. You can just host your own server.

React.useEffect(() => {
  const channel = pusher.subscribe('postEvents');

  channel.bind('newReply', (event) => {
    /* ... */
  });

  return () => {
    channel.unbind('newReply');
    pusher.unsubscribe('postEvents');
  };
}, []);

Interacting with the cache

There are two ways that I typically approach passing a socket message along to React Query. Each are useful in different contexts, so we'll cover both and when you might want to use one over the other.

The first is to invalidate the query by queryKey, so that the next time that query is needed it'll fetch fresh data from the server. This approach minimizes the amount of data over the wire via the websocket and allows users to only fetch the data that they need, when they need it.

The second approach is to send the updated data through the websocket and then manually update the queryClient cache with that new data. This approach is better if you expect to have a frequent, small updates or if you expect a large number of concurrent consumers — and want to avoid triggering a thundering herd every time new data is available.

Invalidate Query

When a socket event occurs, we'll use queryClient.invalidateQueries. The nice thing about React Query here is that if the query is actively being used, it'll refetch that data in the background and show the stale data in the UI until the new data is available. If that query is not actively being used, React Query marks it as invalidated and fetches it the next time the user needs it.

Using the same example from above, we'll fill out the newReply event handler.

React.useEffect(() => {
  const channel = pusher.subscribe('postEvents');

  channel.bind('newReply', (postId) => {
    queryClient.invalidateQueries({
      queryKey: ['replies', postId]
    })
  });

  return () => {
    channel.unbind('newReply');
    pusher.unsubscribe('postEvents');
  };
}, []);

Now on the server, you can send a message to the post-replies channel with the postId as the event data:

pusher.trigger('postEvents', 'newReply', postId);

Manually update the cache

If you expect to have frequent, small updates — like a chat feature — or a large number of concurrent users it is probably better to ship the updated data over the websocket and update the cache by queryKey directly using queryClient.setQueryData.

React.useEffect(() => {
  const channel = pusher.subscribe('postEvents');

  channel.bind('newReply', (reply) => {
    queryClient.setQueryData(
      ['replies', reply.postId],
      (cacheData) => [...cacheData, reply]
    );
  });

  return () => {
    channel.unbind('newReply');
    pusher.unsubscribe('postEvents');
  };
}, []);

In this case, we're still updating the same record by queryKey but instead of telling React Query to fetch the data next time it needs it, we just take the cached value as a parameter in our update function and append our new comment entry to the array.

And on the server side, rather than just sending the postId we'll send the entire reply object.

pusher.trigger('postEvents', 'newReply', reply);