Query Key Factory 구현을 통한 TanStack queryKey 관리하기

배경

회사에서 부분적으로 사용하던 Tanstack Query를 팀 회의를 거쳐 전면 도입하게 되면서, 사용하는 query가 점점 많아지다 보니 좀 더 구조화를 잘 할 수 없을지 고민이 되었다. 그러다 팀원 분을 통해 Query Key Factory라는 라이브러리를 알게 되었는데, 쉬운 방식으로 queryKey 뿐만 아니라 queryFn 및 다른 옵션들까지 함께 관리할 수 있다는 것이 무척 마음에 들어서 도입했었다.

그런데 타입 추론 문제가 있어 해당 라이브러리를 포기하게 되었지만, query key factory의 인터페이스는 포기할 수 없었다. 그래서 직접 구현하여 지금까지 사용해 오고 있다 😎

구현 과정

우리가 원했던 인터페이스는 아래와 비슷했다.

export const users = createQueryKeys('users', {
  detail: (userId: string) => ({
    queryKey: [userId],
    queryFn: () => api.getUser(userId),
  }),
});

// => createQueryKeys output:
// {
//   _def: ['users'],
//   detail: (userId: string) => {
//     queryKey: ['users', 'detail', userId],
//     queryFn: (ctx: QueryFunctionContext) => api.getUser(userId),
//   },
// }

export function useUserDetail(id: string) {
  return useQuery(users.detail(id));
};

라이브러리 예시인데, useQuery(users.detail(id))처럼 사용했을 때 자동으로 unique한 queryKey와 함께 queryFn을 설정할 수 있도록 해서 react query를 사용 하는 데에 드는 시간을 줄이고 싶었다!

그럼 구현 과정을 살펴 보자. 함수명은 query key factory에서 쓰인 것과 동일하게 가져갔다.

• createQueryKeys

queryKeys 객체를 생성하는 함수이다. 아래처럼 사용된다.

export const domainAQueryKeys = createQueryKeys({
  getSomeAPI: (id: number) => ({
    queryKey: [id],
    queryFn: () => getSomeAPI(id),
  }),
  postSomeAPI: {
    queryKey: [],
    queryFn: (body: PostSomeAPIBody) => postSomeAPI(body),
  },
});

• mergeQueryKeys

여러 queryKeys 객체를 merge하는 함수이다. 아래처럼 사용된다.

export const queryKeys = mergeQueryKeys({
  domainA: domainAQueryKeys,
  domainB: domainBQueryKeys,
});

결과적으로 queryKeys.domainA.getSomeAPI(1)을 했을 때 queryKey와 queryFn을 함께 설정할 수 있고, queryKeys.domainA.getSomeAPI(1).queryKey는 ['domainA', 'getSomeAPI', 1]이 되도록 할 것이다.

Types

import type { QueryKey } from '@tanstack/query-core';

export type MergedQueryKeys = Record<string, QueryKeys>;

export type QueryKeys = Record<
  string,
  RequiredQueryOptions | ((...args: any[]) => RequiredQueryOptions)
>;

export interface RequiredQueryOptions {
  queryKey: QueryKey;
  queryFn: (...args: any[]) => Promise;
}

QueryKeys는 RequiredQueryOptions 또는 RequiredQueryOptions를 반환하는 함수를 값으로 갖는 Record이다. 그리고 MergedQueryKeys는 QueryKeys를 값으로 갖는 Record이다.

createQueryKeys

export const createQueryKeys = <T extends QueryKeys>(queryKeys: T) => {
  return queryKeys;
};

이 함수는 로직적으로는 필요하지 않지만, queryKey/queryFn을 등록할 때에는 QueryKeys type을 따르도록 강제하고, 추후 useQuery 등 사용 시 queryFn의 리턴 값이 타입 추론 되도록 하기 위해 필요한 함수이다. extends 문법을 통해 T type의 값을 반환한다.

mergeQueryKeys

export const mergeQueryKeys = <T extends MergedQueryKeys>(allQueryKeys: T) => {
  for (const [domain, queryKeys] of Object.entries(allQueryKeys)) {
    for (const [api, previousValue] of Object.entries(queryKeys)) {
      if (typeof previousValue === 'object') {
        // 객체일 때의 queryKey 변경
        const previousQueryKey = previousValue.queryKey;
        queryKeys[api] = {
          ...previousValue,
          queryKey: [domain, api, ...previousQueryKey],
        };
      } else if (typeof previousValue === 'function') {
        // 함수일 때의 queryKey 변경
        queryKeys[api] = (...args: any[]) => {
          const previousQueryOptions = previousValue(...args);
          const previousQueryKey = previousQueryOptions.queryKey;
          return {
            ...previousQueryOptions,
            queryKey: [domain, api, ...previousQueryKey],
          };
        };
      } else {
        throw new Error(`Invalid format for queryKeys.${domain}.${api}`);
      }
    }
  }

  return allQueryKeys;
};

createQueryKeys와 비슷한 이유로 extends 문법을 사용한다. 그리고 Object.entries() 문법을 통해 unique한 queryKey를 만들어 주는 동작을 수행한다.
이때, 반환되는 T type이 새로운 queryKey 값을 반영하진 못 하는 문제가 있다 😢 하지만 아직 정확한 queryKey 값을 추론해야 하는 상황이 발생하지 않아서 일단 문제를 보류하였다.

결과

queryKeys를 아래와 같이 편리하게 구조화하고 활용할 수 있다 ✨

  const { data, fetchNextPage: fetchNextLogs } = useInfiniteQuery({
    ...queryKeys.accounts.getUserPointsAPI,
    getNextPageParam: lastPage => {
      const next = lastPage.data.next;
      return next ?? undefined;
    },
  });

배운 점

원래 createQueryKeys라는 함수 대신, TypeScript의 satisfies 문법을 사용하여 queryFn의 타입 추론이 가능하도록 만들고 싶었다. 그런데 Babel 문제로 추정되는 문제가 끝까지 해결되지 않아서 결국 사용을 보류하게 되었고, 머리를 싸매고 떠올린 대안이 extends 키워드를 사용하는 것이었다 🥹 extends는 평소에도 알고 있던 문법이지만, 타입 추론에 이게 그렇게 큰 역할을 할 줄이야..!

또한, query key factory를 구현하며 react query 관련 네이밍 컨벤션도 정하게 되어서 네이밍에 쓰던 불필요한 시간도 단축할 수 있었다. 좋은 컨벤션은 높은 효율을 가져오기에, 팀에 이러한 컨벤션들을 쌓아 나가며 생산성을 더욱 높이고 싶은 마음이다.