In part 1, we talked about how single flight mutations allow you to update data, and re-fetch all the relevant updated data for the UI, all in just one roundtrip across the network.
We implemented a trivial solution for this, which is to say that we threw caution (and coupling) to the wind and just re-fetched what we needed in the server function we had for updating data. This worked fine, but it was hardly scalable, or flexible.
Article Series
In this post, we’ll accomplish the same thing, but in a much more flexible way. We’ll define some refetching middleware that we can attach to any server function. The middleware will allow us to specify, via react-query keys, what data we want re-fetched, and it’ll handle everything from there.
We’ll start simple, and keep on adding features and flexibility. Things will get a bit complex by the end, but please don’t think you need to use everything we’ll talk about. In fact, for the vast majority of apps, single flight mutations probably won’t matter at all. And don’t be fooled: simply re-fetching some data in a server function might be good enough for a lot of smaller apps as well.
But in going through all of this we’ll get to see some really cool TanStack, and even TypeScript features. Even if you never use what we go over for single flight mutations, there’s a good chance this content will come in handy for something else.
Our First Middleware
TanStack Query (which we sometimes refer to as react-query, it’s package name) already has a wonderful system of hierarchical keys. Wouldn’t it be great if we could just have our middleware receive the query keys of what we want to refetch, and have it just… work? Have the middleware figure out how to refetch does seem tricky, at first. Sure, our queries have all been simple calls (by design) to server functions. But we can’t pass a server function reference up to the server; functions are not serializable. How could they be? You can send strings and numbers (and booleans) across the wire, serialized as JSON, but sending a function (which can have state, close over context, etc) makes no sense.
Unless they’re TanStack Start server functions, that is.
It turns out the incredible engineers behind this project customized their serialization engine to support server functions. That means you can send a server function to the server, from the client, and it will work fine. Under the covers, server functions have an internal ID. TanStack picks this up, sends the ID, and then de-serializes the ID on the other end.
To make this even easier, why don’t we just attach the server function (and the argument it takes) right in to the query options we already have defined. Then our middleware can take the query keys we want re-fetched, look up the query from TanStack Query internals (which we’ll dive into) and just make everything work.
Let’s Get Started
First we’ll import some goodies
import { createMiddleware, getRouterInstance } from "@tanstack/react-start";
import { QueryClient, QueryKey } from "@tanstack/react-query";Code language: TypeScript (typescript)Next, let’s update our query options for our epics list query (the main list of epics).
export const epicsQueryOptions = (page: number) => {
return queryOptions({
queryKey: ["epics", "list", page],
queryFn: async () => {
const result = await getEpicsList({ data: page });
return result;
},
staleTime: 1000 * 60 * 5,
gcTime: 1000 * 60 * 5,
meta: {
__revalidate: {
serverFn: getEpicsList,
arg: page,
},
},
});
};Code language: TypeScript (typescript)Note the new meta section. This allows us to add any random metadata that we want to our query. Here we send over a reference to the getEpicsList server function, and the arg it takes. If this duplication makes you uneasy, stay tuned. We’ll also update the summary query (for the counts) the same way, though that’s not shown here.
Let’s build this middleware piece by piece.
// the server function and args are all `any`, for now,
// to keep things simple we'll see how to type them in a bit
type RevalidationPayload = {
refetch: {
key: QueryKey;
fn: any;
arg: any;
}[];
};
type RefetchMiddlewareConfig = {
refetch: QueryKey[];
};
export const refetchMiddleware = createMiddleware({ type: "function" })
.inputValidator((config?: RefetchMiddlewareConfig) => config)
.client(async ({ next, data }) => {
const { refetch = [] } = data ?? {};Code language: TypeScript (typescript)We define an input to the middleware. This input will automatically get merged with whatever input is defined on any server function this middleware winds up attached to.
We define our input as optional (config?) since it’s entirely possible we might want to sometimes call our server function and simply not refetch anything.
Now we start our client callback, which runs directly in our browser. We’ll first grab the array of query keys we want refetched.
const { refetch = [] } = data ?? {};Code language: TypeScript (typescript)Then we’ll get our queryClient and the cache attached to it, and define the payload we’ll send to the server callback of our middleware, which will do the actual refetching.
If you’ve never touched TanStack’s middleware before and are feeling overwhelmed, my middleware post might be a good introduction.
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
const cache = queryClient.getQueryCache();
const revalidate: RevalidationPayload = {
refetch: [],
};Code language: TypeScript (typescript)Our queryClient is already attached to the main TanStack router context, so we can get the router, and just grab it.
Remember before when we added that __revalidate payload to our query options, with the server function, and arg? Let’s look in our query cache for each key, and retrieve the query options for them.
refetch.forEach((key: QueryKey) => {
const entry = cache.find({ queryKey: key, exact: true });
if (!entry) return;
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}
});Code language: TypeScript (typescript)The check if (!entry) return; protects us from refetches being requested for queries that don’t exist in cache—ie, if they’ve never been fetched in the UI. If that happens, just skip to the next one. We have no way to refetch it if we don’t have the serverFn.
You could expand the input to this middleware and send up a different payload of query keys, along with the actual refetching payload (including server function and arg) for queries you absolutely want run, even if they haven’t yet been requested. Perhaps you’re planning on redirecting after the mutation, and you want that new page’s data prefetched. We won’t implement that here, but it’s just a variation on this same theme. These pieces are all very composable, so build whatever you happen to need!
This code then grabs the meta object, and puts the properties onto the payload we’ll send to the server.
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}Code language: TypeScript (typescript)Try not to let the various any types bother you; I’m omitting some type definitions that would have been straightforward to define, in order to help prevent this long post from getting even longer.
Calling next triggers the actual invocation of the server function (and any other middleware in the chain). The sendContext arg allows us to send data from the client, up to the server. The server is allowed to call next with a sendContext payload that sends data back to the client.
const result = await next({
sendContext: {
revalidate,
},
});Code language: TypeScript (typescript)The result payload is what comes back from the server function invocation. The context object on it will have a payloads array, returned from the .server callback just below, with entries containing a key (the query key), and result (the actual data). We’ll loop it, and update the query data accordingly.
We’ll fix the TS error covered up with // @ts-expect-error momentarily.
// @ts-expect-error
for (const entry of result.context?.payloads ?? []) {
queryClient.setQueryData(entry.key, entry.result);
}
return result;Code language: TypeScript (typescript)The Server Callback
The server callback looks like this, in its entirety.
.server(async ({ next, context }) => {
const result = await next({
sendContext: {
payloads: [] as any[]
}
});
const allPayloads = context.revalidate.refetch.map(refetchPayload => {
return {
key: refetchPayload.key,
result: refetchPayload.fn({ data: refetchPayload.arg })
};
});
for (const refetchPayload of allPayloads) {
result.sendContext.payloads.push({
key: refetchPayload.key,
result: await refetchPayload.result
});
}
return result;Code language: TypeScript (typescript)We immediately call next(), which runs the actual server function this middleware is attached to. We pass a payloads array in sendContext. This governs what gets sent back to the client callback (that’s how .client got the payloads array we just saw it looping through).
Then we run through the revalidate payloads sent up from the client. The client sent them via sendContext, and we read them from the context object (send context, get it?). We then call all the server functions, and add to that payloads array.
Here’s the entire middleware:
export const refetchMiddleware = createMiddleware({ type: "function" })
.inputValidator((config?: RefetchMiddlewareConfig) => config)
.client(async ({ next, data }) => {
const { refetch = [] } = data ?? {};
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
const cache = queryClient.getQueryCache();
const revalidate: RevalidationPayload = {
refetch: [],
};
refetch.forEach((key: QueryKey) => {
const entry = cache.find({ queryKey: key, exact: true });
if (!entry) return;
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}
});
const result = await next({
sendContext: {
revalidate,
},
});
// @ts-expect-error
for (const entry of result.context?.payloads ?? []) {
queryClient.setQueryData(entry.key, entry.result);
}
return result;
})
.server(async ({ next, context }) => {
const result = await next({
sendContext: {
payloads: [] as any[],
},
});
const allPayloads = context.revalidate.refetch.map(refetchPayload => {
return {
key: refetchPayload.key,
result: refetchPayload.fn({ data: refetchPayload.arg }),
};
});
for (const refetchPayload of allPayloads) {
result.sendContext.payloads.push({
key: refetchPayload.key,
result: await refetchPayload.result,
});
}
return result;
});Code language: TypeScript (typescript)Fixing the TypeScript Error
Why is this line invalid?
// @ts-expect-error
for (const entry of result.context?.payloads ?? []) {Code language: TypeScript (typescript)This line runs in the .client callback, after we call next(). Essentially, we’re trying to read properties sent back to the client, from the server (via the sendContext payload). This runs, and works properly. But why don’t the types line up?
I covered this in my Middleware post linked above, but our server callback can see what gets sent to it from the client, but the reverse is not true. This knowledge just inherently does not go in both directions; the type inference cannot run backwards.
The solution is simple: just break the middleware into two pieces, and make one of them a middleware dependency on the other.
const prelimRefetchMiddleware = createMiddleware({ type: "function" })
.inputValidator((config?: RefetchMiddlewareConfig) => config)
.client(async ({ next, data }) => {
const { refetch = [] } = data ?? {};
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
// same
// as
// before
return await next({
sendContext: {
revalidate,
},
});
// those last few lines are removed
})
.server(async ({ next, context }) => {
const result = await next({
sendContext: {
payloads: [] as any[],
},
});
// exactly the same as before
return result;
});
export const refetchMiddleware = createMiddleware({ type: "function" })
.middleware([prelimRefetchMiddleware]) // <-------- connect them!
.client(async ({ next }) => {
const result = await next();
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
// and here's those last few lines we removed from above
for (const entry of result.context?.payloads ?? []) {
queryClient.setQueryData(entry.key, entry.result);
}
return result;
});Code language: TypeScript (typescript)It’s the same as before, except everything in the .client callback after the call to next() is now in its own middleware. The rest is in a different middleware, which is an input to this one. Now when we call next in refetchMiddleware, TypeScript is able to see the data that’s been sent down from the server, since that was done in prelimRefetchMiddleware, which is an input to this middleware, which allows TypeScript to fully see the flow of types.
Wiring it Up
Now we can take our server function for updating an epic, remove the refecthes, and add our refetch middleware.
export const updateEpic = createServerFn({ method: "POST" })
.middleware([refetchMiddleware])
.inputValidator((obj: { id: number; name: string }) => obj)
.handler(async ({ data }) => {
await new Promise(resolve => setTimeout(resolve, 1000 * Math.random()));
await db.update(epicsTable).set({ name: data.name }).where(eq(epicsTable.id, data.id));
});Code language: TypeScript (typescript)We set it up to call from our React component with the useServerFn hook, which handles things like errors and redirects automatically.
const runSave = useServerFn(updateEpic);Code language: TypeScript (typescript)Remember when I said that inputs to middleware are automatically merged with inputs to the underlying server function? We can see that first hand when we call the server function.
unknown[] is the correct type for react-query query keys)Now we can call it, and specify the queries we want refetched.
await runSave({
data: {
id: epic.id,
name: newValue,
refetch: [
["epics", "list", 1],
["epics", "list", "summary"],
],
},
});Code language: TypeScript (typescript)When we run it, it works. Both the list of epics, and also the summary list correctly update with our changes, without any new requests in the network tab. When testing single flight mutations, we’re not really looking for something to indicate that it worked, but rather a lack of new network requests, for updated data.
Improving Things
Query keys are hierarchical in react-query. You might already be familiar with this. Normally, when updating data, it would be common to do something like:
queryClient.invalidateQueries({ queryKey: ["epics", "list"] });Code language: TypeScript (typescript)Which refetches any queries whose key starts with ["epics", "list"]. Can we do something similar in our middleware? Like, just pass in that key prefix, and have it find, and refetch whatever’s there?
Let’s do it!
Getting the matching keys will be slightly more complicated. Each key we pass up will potentially be a key prefix, matching multiple entries, so we’ll use flatMap to find all matches, with the nifty cache.findAll method.
const allQueriesFound = refetch.flatMap(
k => cache.findAll({ queryKey: k, exact: false })
);Code language: TypeScript (typescript)Now we loop them, and do the same thing as before.
const allQueriesFound = refetch.flatMap(
k => cache.findAll({ queryKey: k, exact: false })
);
allQueriesFound.forEach(entry => {
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key: entry.queryKey,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}
});Code language: TypeScript (typescript)This works!
Going Deeper
Our solution still isn’t ideal. What if we page around in our epics page (up to page 2, up to page 3, then back down). Our solution will find page 1, and our summary query, but also pages 2 and 3, since they’re now in cache. But pages 2 and 3 aren’t really active, and we shouldn’t refetch them, since they’re not even being displayed.
Let’s change our code to only refetch active queries. Detecting whether a query entry is actually active is as simple as adding the type argument to findAll.
cache.findAll({ queryKey: key, exact: false, type: "active" });Code language: TypeScript (typescript)Our code now looks like this:
const allQueriesFound = refetch.flatMap(key => cache.findAll({ queryKey: key, exact: false, type: "active" }));
allQueriesFound.forEach(entry => {
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key: entry.queryKey,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}
});Code language: TypeScript (typescript)Even Deeper
This works. But when you think about it, those other, inactive queries should probably be invalidated. We don’t want to waste resources refetching them immediately, since they’re not being used; but if the user were to browse back to those pages, we probably want the data refetched. Tanstack Query makes that easy, via the invalidateQueries method.
We’ll add this to the client callback of the middleware we feed into.
data?.refetch.forEach(key => {
queryClient.invalidateQueries({ queryKey: key, exact: false, type: "inactive", refetchType: "none" });
});Code language: TypeScript (typescript)Loop the query keys we passed in, and invalidate any of the inactive queries (the active ones have already been refetched), but without refetching them.
Here’s our entire, updated middleware.
const prelimRefetchMiddleware = createMiddleware({ type: "function" })
.inputValidator((config?: RefetchMiddlewareConfig) => config)
.client(async ({ next, data }) => {
const { refetch = [] } = data ?? {};
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
const cache = queryClient.getQueryCache();
const revalidate: RevalidationPayload = {
refetch: [],
};
const allQueriesFound = refetch.flatMap(key => cache.findAll({ queryKey: key, exact: false, type: "active" }));
allQueriesFound.forEach(entry => {
const revalidatePayload: any = entry?.meta?.__revalidate ?? null;
if (revalidatePayload) {
revalidate.refetch.push({
key: entry.queryKey,
fn: revalidatePayload.serverFn,
arg: revalidatePayload.arg,
});
}
});
return await next({
sendContext: {
revalidate,
},
});
})
.server(async ({ next, context }) => {
const result = await next({
sendContext: {
payloads: [] as any[],
},
});
const allPayloads = context.revalidate.refetch.map(refetchPayload => {
return {
key: refetchPayload.key,
result: refetchPayload.fn({ data: refetchPayload.arg }),
};
});
for (const refetchPayload of allPayloads) {
result.sendContext.payloads.push({
key: refetchPayload.key,
result: await refetchPayload.result,
});
}
return result;
});
export const refetchMiddleware = createMiddleware({ type: "function" })
.middleware([prelimRefetchMiddleware])
.client(async ({ data, next }) => {
const result = await next();
const router = await getRouterInstance();
const queryClient: QueryClient = router.options.context.queryClient;
for (const entry of result.context?.payloads ?? []) {
queryClient.setQueryData(entry.key, entry.result, { updatedAt: Date.now() });
}
data?.refetch.forEach(key => {
queryClient.invalidateQueries({ queryKey: key, exact: false, type: "inactive", refetchType: "none" });
});
return result;
});Code language: TypeScript (typescript)We tell TanStack Query to invalidate (but not refetch) any inactive queries matching our key.
This works perfectly. If we browse up to pages 2 and 3, and then back to page 1, then edit a todo, we do in fact see our list, and summary list update. If we then page back to page 2, and 3, we’ll see network requests fire to get fresh data.
Icing on the Cake
Remember when we added the server function, and the arg it takes to our query options?
export const epicsQueryOptions = (page: number) => {
return queryOptions({
queryKey: ["epics", "list", page],
queryFn: async () => {
const result = await getEpicsList({ data: page });
return result;
},
staleTime: 1000 * 60 * 5,
gcTime: 1000 * 60 * 5,
meta: {
__revalidate: {
serverFn: getEpicsList,
arg: page,
},
},
});
};Code language: TypeScript (typescript)I briefly noted that it was a bit gross to duplicate the server function, and arg in both our meta object, and our queryFn. Let’s fix this.
Let’s start with the simplest possible helper to remove this duplication, and as before, iterate on it.
export function refetchedQueryOptions(queryKey: QueryKey, serverFn: any, arg?: any) {
const queryKeyToUse = [...queryKey];
if (arg != null) {
queryKeyToUse.push(arg);
}
return queryOptions({
queryKey: queryKeyToUse,
queryFn: async () => {
return serverFn({ data: arg });
},
meta: {
__revalidate: {
serverFn,
arg,
},
},
});
}Code language: TypeScript (typescript)It’s just a simple helper that takes in your query key, server function and arg, and returns back some of our query options: our queryKey (to which we add whatever argument we need for the server function), the queryFn which calls the server function, and our meta object.
Our epics list query now looks like this:
export const epicsQueryOptions = (page: number) => {
return queryOptions({
...refetchedQueryOptions(["epics", "list"], getEpicsList, page),
staleTime: 1000 * 60 * 5,
gcTime: 1000 * 60 * 5,
});
};Code language: TypeScript (typescript)This works, but it’s not great. We have any types everywhere, which means the argument we pass to our server function is never type checked. Even worse, the return value of our queryFn is not type checked, which means our queries (like this very epics list query) now return any.
Let’s add some typings. Server functions are just functions. They take a single object argument, and if the server function has defined an input, then that argument will have a data property for that argument. That’s a lot of words to say what we already know. When we call a server function, we pass our argument like this
const result = await runSaveSimple({
data: {
id: epic.id,
name: newValue,
},
});Code language: TypeScript (typescript)How’s this for a second draft?
export function refetchedQueryOptions<T extends (arg: { data: any }) => Promise<any>>(
queryKey: QueryKey,
serverFn: T,
arg: Parameters<T>[0]["data"],
) {
const queryKeyToUse = [...queryKey];
if (arg != null) {
queryKeyToUse.push(arg);
}
return queryOptions({
queryKey: queryKeyToUse,
queryFn: async (): Promise<Awaited<ReturnType<T>>> => {
return serverFn({ data: arg });
},
meta: {
__revalidate: {
serverFn,
arg,
},
},
});
}Code language: TypeScript (typescript)We’ve constrained our server function to an async function which takes a data prop on its object arg, and we’ve used that to statically type the argument. This is good, but we get an error when we use this on server functions which have no arguments.
...refetchedQueryOptions(["epics", "list", "summary"], getEpicsSummary)
// Expected 3 arguments, but got 2.Code language: TypeScript (typescript)Adding an undefined does fix this, and everything works.
...refetchedQueryOptions(["epics", "list", "summary"], getEpicsSummary, undefined),Code language: TypeScript (typescript)If you’re normal, you’re probably happy with that. And you should be. But if you’re weird like me, you might wonder if you can’t make it perfect. Ideally it would be cool if we could pass a statically typed argument when using a server function that takes an input, and when using a server function with no input, pass nothing.
TypeScript has a feature exactly for this: overloaded functions.
This post is already far too long, so I’ll post the code, and leave deciphering it as an exercise for the reader (and likely a future blog post).
import { QueryKey, queryOptions } from "@tanstack/react-query";
type AnyAsyncFn = (...args: any[]) => Promise<any>;
type ServerFnArgs<TFn extends AnyAsyncFn> = Parameters<TFn>[0] extends infer TRootArgs
? TRootArgs extends { data: infer TResult }
? TResult
: undefined
: never;
type ServerFnHasArgs<TFn extends AnyAsyncFn> = ServerFnArgs<TFn> extends infer U ? (U extends undefined ? false : true) : false;
type ServerFnWithArgs<TFn extends AnyAsyncFn> = ServerFnHasArgs<TFn> extends true ? TFn : never;
type ServerFnWithoutArgs<TFn extends AnyAsyncFn> = ServerFnHasArgs<TFn> extends false ? TFn : never;
type RefetchQueryOptions<T> = {
queryKey: QueryKey;
queryFn?: (_: any) => Promise<T>;
meta?: any;
};
type ValidateServerFunction<Provided, Expected> = Provided extends Expected ? Provided : "This server function requires an argument!";
export function refetchedQueryOptions<TFn extends AnyAsyncFn>(
queryKey: QueryKey,
serverFn: ServerFnWithArgs<TFn>,
arg: Parameters<TFn>[0]["data"],
): RefetchQueryOptions<Awaited<ReturnType<TFn>>>;
export function refetchedQueryOptions<TFn extends AnyAsyncFn>(
queryKey: QueryKey,
serverFn: ValidateServerFunction<TFn, ServerFnWithoutArgs<TFn>>,
): RefetchQueryOptions<Awaited<ReturnType<TFn>>>;
export function refetchedQueryOptions<TFn extends AnyAsyncFn>(
queryKey: QueryKey,
serverFn: ServerFnWithoutArgs<TFn> | ServerFnWithArgs<TFn>,
arg?: Parameters<TFn>[0]["data"],
): RefetchQueryOptions<Awaited<ReturnType<TFn>>> {
const queryKeyToUse = [...queryKey];
if (arg != null) {
queryKeyToUse.push(arg);
}
return queryOptions({
queryKey: queryKeyToUse,
queryFn: async () => {
return serverFn({ data: arg });
},
meta: {
__revalidate: {
serverFn,
arg,
},
},
});
}Code language: TypeScript (typescript)With this in place we can now call it with server functions that take an argument.
export const epicsQueryOptions = (page: number) => {
return queryOptions({
...refetchedQueryOptions(["epics", "list"], getEpicsList, page),
staleTime: 1000 * 60 * 5,
gcTime: 1000 * 60 * 5,
});
};Code language: TypeScript (typescript)The parameter is now checked. It errors with the wrong type.
...refetchedQueryOptions(["epics", "list"], getEpicsList, "")
// Argument of type 'string' is not assignable to parameter of type 'number'.Code language: TypeScript (typescript)It errors if you pass no argument as well.
...refetchedQueryOptions(["epics", "list"], getEpicsList)
// Argument of type 'RequiredFetcher<undefined, (page: number) => number, Promise<{ id: number; name: string; }[]>>' is not assignable to parameter of type '"This server function requires an argument!"'.Code language: TypeScript (typescript)That last error isn’t the clearest, but if you read to the end you get a pretty solid hint as to what’s wrong, thanks to this dandy little helper.
type ValidateServerFunction<Provided, Expected> = Provided extends Expected ? Provided : "This server function requires an argument!";Code language: TypeScript (typescript)And it works with a server function that takes no arguments. Again, a full explanation of this TypeScript will have to wait for a future post.
Concluding Thoughts
Single flight mutations are a great tool for speeding up updates within your web app, particularly when it’s a performance boost to avoid follow-up network requests for data after an initial mutation. Hopefully this post has shown you the tools needed to put it all together.
