Enhance your tanstack/query keys with Proxy

Original post

Treat it as an extension of the iconic blog post Effective react query keys. If you already use a similar structure (scope, keys, hooks) and have one shared context like shop in every key, this post is for you.

Managing query keys

As your application grows, you have to think about some design decisions. One of them can be managing your query keys in Tanstack Query. Recently I encountered this situation at work. I had a facade for all react-query hooks and used the strategy from Effective react query keys, so I felt pretty safe with my design decisions.

It looked something like this:

export function useQueryWithLocation<
  TQueryFnData = unknown,
  TErrorType = ErrorType,
  TSelectData = TQueryFnData,
>(
  queryKey: ((shop: string) => QueryKey) | QueryKey,
  queryFn: (
    shop: string,
    context: QueryFunctionContext,
  ) => Promise<TQueryFnData>,
  options: Omit<
    UseQueryOptions<TQueryFnData, TErrorType, TSelectData>,
    "queryKey" | "queryFn"
  > &
    StatusHandlers<TSelectData, TErrorType> = {},
) {
  const shop = useShop();

  const queryResult = useQuery<TQueryFnData, TErrorType, TSelectData>(
    typeof queryKey === "function"
      ? queryKey(shop)
      : Array.isArray(queryKey)
        ? [shop, ...queryKey]
        : [shop, queryKey],
    (context) => queryFn(shop, context),
    options,
  );

  return queryResult;
}

export const searchScopes = {
  searchResults: 'searchResults',
};

export const searchKeys = {
  searchResults: (shop: string, orderId: string, searchValue: string) => [
    searchScopes.searchResults,
    shop,
    orderId,
    searchValue,
  ],
};

export const useSearchResults = ({
  options,
  orderId,
  searchValue,
}: CustomUseQueryOptions<SearchDTO, ErrorType> & {
  orderId: string;
  searchValue: string;
}) => {
  return useQueryWithLocation<SearchDTO, ErrorType>(
    (shop: string) => searchKeys.searchResults(shop, orderId, searchValue),
    (shop: string, { signal }) =>
      getFoundItems({
        shop,
        orderId,
        finderOptions: { query: searchValue },
        signal,
      }),
    options
  );
};

In this setup, I had a facade named useQueryWithLocation which added the current shop to all queries. I also had domain-scoped queries for every functionality where I had a scope, keys, and query hooks which used useQueryWithLocation and defined keys. A structure like this is great when you have to invalidate a query or do something with it. You don't have to manually define keys; you can use your const map of keys. For example:

await queryClient.invalidateQueries({
  queryKey: searchKeys.searchResults(shop, orderId),
  exact: false,
});

This can invalidate all search queries for a given order and shop. But if you have read Effective react query keys, you already know that. So one question pops up...

Great, so what's the problem?

The problem started when we implemented roles into our application. In the previous code samples, I passed the shop value. It is passed in almost every query because most queries are scoped by shop. Now we have to add a new value to every query, which is role. I started thinking about how to do that without much work. And it was hard. I was thinking about setting the role in local storage. That was the easy way but not the correct one. I was leaning toward the harder way, which was adding the role in a similar fashion to how shop was added. But that wasn't ideal. If I did that, I would have to add it to useQueryWithLocation and then manually to all key objects and fetch functions (in this example, getFoundItems). In the AI world that could probably be automated, but if something this simple is so hard to change, I started thinking about refactoring because clearly the current solution doesn't scale well.

What changes were needed?

I started by renaming the default useQuery facade to useEnhancedQuery and changing some small details.

export function useEnhancedQuery<
  TQueryFnData = unknown,
  TErrorType = ErrorType,
  TSelectData = TQueryFnData,
>({
  endpoint,
  queryKey,
  options = {},
  fetchOptions,
}: {
  endpoint: (params: { shop: string }) => string;
  options?: Omit<
    UseQueryOptions<TQueryFnData, TErrorType, TSelectData>,
    "queryKey" | "queryFn"
  > &
    StatusHandlers<TSelectData, TErrorType>;
  fetchOptions: Omit<ConfigType, "shop" | "role">;
  queryKey: QueryKey;
}) {
  const shop = useActiveShop();
  const role = useRole();
  const enhancedQueryKey = useEnhancedQueryKey(queryKey);

  return useQuery(
    enhancedQueryKey,
    (context) =>
      client(endpoint({ shop }), {
        ...fetchOptions,
        signal: context.signal,
        shop,
        role,
      }),
    options,
  );
}

Notice that I added the role there. There is also a new hook called useEnhancedQueryKey. It's very simple, and I will show it soon. I also changed the way of passing the fetch function. Now useEnhancedQuery accepts an endpoint (which is a callback with the shop passed and returning a string). I did that because it's now easier to pass new parameters straight away to the client (fetch wrapper) function. Now let's see useEnhancedQueryKey.

export const useEnhancedQueryKey = (queryKey: QueryKey) => {
  const shop = useActiveShop();
  const role = useRole();

  return useMemo(
    () => [shop, role, ...queryKey],
    [queryKey, shop, role]
  );
};

In short, it only adds the shop and context role at the start of the query key that was passed. Next, we can see the changed usage of useSearchResults.

export const searchScopes = {
  searchResults: "searchResults",
};

const searchKeys = {
  searchResults: (orderId: string, searchValue: string) => [
    searchScopes.searchResults,
    orderId,
    searchValue,
  ],
};

Search scopes remain unchanged. There are changes in searchKeys, though. The first small detail is that it's no longer exported. The second is that there is no shop being passed, and for the same reason there is no new role attribute. As you know from the previous sample, useEnhancedQueryKey is now responsible for that. But what about the invalidate-query example? Now when we want to invalidate, we can't use search keys because they are not exported. We have a solution for that:

export const useEnhancedSearchKeys = () => {
  return useEnhanceQueryKeys(searchKeys);
};

Hm, that doesn't explain much... This hook is now responsible for returning the search keys. It is more restrictive because we can't use hooks in regular functions outside React, so that's not ideal. But I think there are more positives to that approach. Let's see the details of the useEnhanceQueryKeys hook:

export const useEnhanceQueryKeys = <T extends Record<string, (...args: any[]) => QueryKey>>(
  queryKeys: T
) => {
  const enhanceQueryKey = useEnhanceQueryKey();

  return useMemo(() => {
    return new Proxy(queryKeys, {
      get(target, prop, receiver) {
        return (...props: Parameters<typeof Reflect.get<typeof target, typeof prop>>) => {
          return enhanceQueryKey(Reflect.get(target, prop, receiver)(...props));
        };
      },
    });
  }, [enhanceQueryKey, queryKeys]);
};

Now that's more interesting! I don't know about you, but this was the first time I used Proxy. I have known about this object for a long time but never found a suitable use case. Now I have it (at least I think I have 😅). The enhanceQueryKey function is very similar to the useEnhancedQueryKeys hook, but it's a function, not a regular array.

export const useEnhanceQueryKey = () => {
  const enhancedKeys = useEnhancedQueryKey([]);
  return (queryKey: QueryKey) => [...enhancedKeys, ...queryKey];
};

We use that function to enhance regular query keys with fixed attributes—in our case, shop and role. A proxy works in a way that it wraps typical actions (for example, get) on that object. In our case, we want to add our shop and role to the regular return value from our object. Let's see the flow:

const searchKeys = useEnhancedSearchKeys();

searchKeys.searchResults(orderId, "example search");
// this returns ['shop-1234','admin',orderId,'example search']

// in our proxy we can replace
// enhanceQueryKey(Reflect.get(target, prop, receiver)(...props));
// with
// enhanceQueryKey(searchKeys.searchResults(orderId, "example search"));

// so now when we want to invalidate every search query for an order we can pass these search keys like this:

await queryClient.invalidateQueries({
  queryKey: searchKeys.searchResults(orderId), // shop and role are included
  exact: false,
});

In this approach, when in the future I have to add a new attribute to all queries, I have an easy way to extend it. I only have to change it in useEnhanceQueryKey and useEnhancedQuery (and the client if it's needed in the fetch call). In our case, I also have to add two headers for every fetch call:

X-Context-Store and X-Context-Role

I wasn't a big fan of it because I don't see much sense in it, but it is for a preauth check as I understand it. So it's not the only source of auth checking (because that's unsafe), but it helps with rejecting some calls before they reach the later stages of auth. But now, after the refactor, adding these two headers is relatively easy, as you saw in the useEnhancedQuery example when I pass role and shop to the client.

The last chunk of code is the usage of useSearchResults with the new useEnhancedQuery call.

export const useSearchResults = ({
  options,
  orderId,
  searchValue,
}: CustomUseQueryOptions<SearchDTO, ErrorType> & {
  orderId: string;
  searchValue: string;
}) => {
  return useEnhancedQuery<SearchDTO, ErrorType>({
    queryKey: searchKeys.searchResults(orderId, searchValue),
    endpoint: ({ shop }) =>
      `search/${shop}/orders/${orderId}/items${createQueryString({ query: searchValue })}`,
    fetchOptions: {
      schema: getFoundItemsSchema,
      errorObject: searchError,
      transactionId: GET_FOUND_ITEMS,
    },
    ...options,
  });
};

Here we pass the query key similar to the previous example, but now it's enhanced automatically. As I mentioned before, we now pass the endpoint as a string and a fetch-options config (there is a schema for validation, transactionId, and an error object for monitoring). The usage is very similar to the previous case, but now it's much more scalable.

Conclusion

One of the findings of this is that you shouldn't be scared of refactoring sometimes. If the time for adding a new thing is close to that of refactoring it into something better, don't even hesitate. Even if it's 1 or 2 days more for a simple task, I think it's worth it because next time you will be prepared and the next change will take you a few minutes. The next thing is that some general knowledge is worth it even if you don't necessarily use it right away. I know you can handle this case without using a proxy, but for me it was a very satisfying moment when I could use it and remembered that in that case it would be suitable.