Improve TypeScript support for dynamic queryKey with union types

[onyedikachi23]

Describe the bug

Indexing a query factory with a union type (UserType = "company" | "club" | "player" | "supporter") causes TypeScript type mismatches in useQuery, ensureQueryData, getQueryData, and so on. The issue arises because entityProfileQueries[userType] returns a union of query options with distinct queryKey types (e.g., { userType: "player" } vs. { userType: "supporter" }), while these utilities expect a single type.

The player and supporter queries share ServerPlayerResponse and a select function, but the error occurs for all UserType values due to the queryKey union.

Your minimal, reproducible example

https://stackblitz.com/edit/vitejs-vite-yzi9yarp?file=src%2Floader.ts

Steps to reproduce

In the MRE, the error arises from this line in src/loader.ts:

const data = await queryClient.ensureQueryData(queryOptions);

queryOptions is underlined with an error, and typeof data is incorrectly resolving to ServerPlayerResponse.

Expected behavior

The returned data from queryClient.ensureQueryData(queryOptions) should be a union of:

type TData = CompanyProfileData | ClubProfileData | PlayerProfileData

Noting that userType - player and supporter, share exactly the same TQueryFnData and TData from the select option, but queryKey differ by { userType: "player" } vs. { userType: "supporter" }

How often does this bug happen?

Every time

Screenshots or Videos

No response

Platform

• OS: all
• browser: all
• version: latest

Tanstack Query adapter

react-query

TanStack Query version

latest

TypeScript version

latest

Additional context

Affects useQuery, ensureQueryData, getQueryData, and so on.

[onyedikachi23]

Please an immediate possible workaround is what I need for now.

[cseas]

Please an immediate possible workaround is what I need for now.

You can workaround this by providing an explicit type cast.

From your reproducible example:

import type { EnsureQueryDataOptions } from '@tanstack/react-query';
...
const data = await queryClient.ensureQueryData(queryOptions as EnsureQueryDataOptions);

For figuring out where the type mismatch is happening, you might need to add explicit type definitions to the entityProfileQueries object and work your way backwards to pinpoint the type mismatch.

The all: () => ['entity-profile'] as const, part of the object is likely making the inferred type of the object complicated as its different than the rest of the object entries.

[onyedikachi23]

as EnsureQueryDataOptions will remove the original correct types. Also, having to write explicit type definitions for that query factory will introduce complexities in managing the types.

[chekrd]

We built a robust type layer over RQ for this reason years ago. It's not an open source repository, though I can publish the browser part. The idea is:

Having robust key definiton for each query. "Query keys" are "leafs" - the runtime cache key shapes. "Filter keys" are abstract - usefull for quering from the cache, cannot exist at runtime as those do not represent a specific entity/cache value.

export const contentTypeKeys = {
  filterKeys: {
    all: [{ entity: 'contentTypes' }] as const,
    allForEnvironment: (projectEnvironmentId: Uuid) =>
      [{ ...contentTypeKeys.filterKeys.all[0], projectEnvironmentId }] as const,
    detailsForEnvironment: (projectEnvironmentId: Uuid) =>
      [
        {
          ...contentTypeKeys.filterKeys.allForEnvironment(projectEnvironmentId)[0],
          scope: 'detail',
        },
      ] as const,
    listsForEnvironment: (projectEnvironmentId: Uuid) =>
      [
        { ...contentTypeKeys.filterKeys.allForEnvironment(projectEnvironmentId)[0], scope: 'list' },
      ] as const,
  },
  queryKeys: {
    detail: (projectEnvironmentId: Uuid, detailId: Uuid) =>
      [
        { ...contentTypeKeys.filterKeys.detailsForEnvironment(projectEnvironmentId)[0], detailId },
      ] as const,
    list: (projectEnvironmentId: Uuid, filter?: ReadonlyRecord<string, string>) =>
      [
        {
          ...contentTypeKeys.filterKeys.listsForEnvironment(projectEnvironmentId)[0],
          filter: filter ?? {},
        },
      ] as const,
  },
} as const;

type QueryKeyNameToCacheModel = {
  readonly detail: ContentTypeServerModel;
  readonly list: ContentTypesServerModel;
};

Offtopic: The filter: filter ?? {}, is a workaround for veery old RQ bug #3741 that I discovered during the implementation of our type layer - still not fixed at the time of writing this comment, still thinking that it is a serious issue in robust systems.

Then we have this brutal thing:

export type ContentTypeKeyMap = QueryKeyMap<typeof contentTypeKeys, QueryKeyNameToCacheModel>;

QueryKeyMap<> returns detailed relations between key type and value types. We use the resulting type to derive key/value combinations from it. Like this one, in contentType query/mutation definition:

type QueryKey = ContentTypeKeyMap['detail']['queryKey'];
type CacheModel = CacheModelForKey<QueryKey, ContentTypeKeyMap>;

Not a real code but to make it short, the runtime query definition:

const createQueryKey = (params: Params): QueryKey =>
  contentTypeKeys.queryKeys.detail(params.projectEnvironmentId, params.detailId);

const queryFn: QueryFn<QueryKey, CacheModel> = ({
  queryKey: [{ projectEnvironmentId, detailId }],
}) => contentTypeRepository.getContentType(projectEnvironmentId, detailId);

const useContentTypeQuery = (projectEnvironmentId: Uuid) => useQuery({ queryFn, queryKey: createQueryKey({ projectEnvironmentId }) });

The ugly part is registering the client for each entity:

export const contentTypesClient = createEntityQueryClient<ContentTypeKeyMap, 'contentTypesClient'>(
  'contentTypesClient',
);

export const queryClients = (queryClient: QueryClient) => ({
  ...contentTypesClient(queryClient),
  ...anotherEntityClient(queryClient),
});

Thats because we need to modify the RQ client type:

/**
 * QueryClient creator per entity. Decorates QueryClient with strict types.
 */
export const createEntityQueryClient =
  <TEntityCacheKeyMap extends EntityQueryKeyMap, TEntityName extends `${string}Client`>(
    entityName: TEntityName,
  ) =>
  (queryClient: QueryClient) => {
    const _queryClient = queryClient as OriginQueryClient;

    const decoratedMethods = {
      getQueryState: <TQueryKey extends QueryKeysUnion<TEntityCacheKeyMap>>(queryKey: TQueryKey) =>
        _queryClient.getQueryState<CacheModelForKey<TQueryKey, TEntityCacheKeyMap>>(queryKey),
      getQueryData: <TQueryKey extends QueryKeysUnion<TEntityCacheKeyMap>>(
        queryKey: TQueryKey,
      ): QueryKeyToCacheModelTuplesForKey<TQueryKey, TEntityCacheKeyMap> | undefined =>
        _queryClient.getQueryData(queryKey),
      getQueriesData: <TAnyKey extends AllKeysUnion<TEntityCacheKeyMap>>(
        filters: QueryFilters<TAnyKey>,
      ): ReadonlyArray<QueryKeyToCacheModelTuplesForKey<TAnyKey, TEntityCacheKeyMap>> =>
        _queryClient.getQueriesData(filters) as any,
      setQueryData: <TQueryKey extends QueryKeysUnion<TEntityCacheKeyMap>>(
        queryKey: TQueryKey,
        updater: QueryDataUpdater<CacheModelForKey<TQueryKey, TEntityCacheKeyMap>>,
        options?: SetDataOptions,
      ): KeyToCacheModelOrUndefinedTuple<
        QueryKeyToCacheModelTuplesForKey<TQueryKey, TEntityCacheKeyMap>
      > =>
        _queryClient.setQueryData(queryKey, updater, options) as KeyToCacheModelOrUndefinedTuple<
          QueryKeyToCacheModelTuplesForKey<TQueryKey, TEntityCacheKeyMap>
        >,
      setQueriesData: <TAnyKey extends AllKeysUnion<TEntityCacheKeyMap>>(
        filters: QueryFilters<TAnyKey>,
        updater: QueryDataUpdater<CacheModelForKey<TAnyKey, TEntityCacheKeyMap>>,
        options?: SetDataOptions,
      ): ReadonlyArray<
        KeyToCacheModelOrUndefinedTuple<
          QueryKeyToCacheModelTuplesForKey<TAnyKey, TEntityCacheKeyMap>
        >
      > => _queryClient.setQueriesData(filters, updater, options) as any,
    };

    return createObjectWithExplicitSingleProperty(entityName, decoratedMethods);
  };

/**
 * Force explicit property name string on a type level.
 * TypeScript is not able to derive property name string from the function argument.
 */
const createObjectWithExplicitSingleProperty = <TName extends string, TValue>(
  name: TName,
  value: TValue,
): ReadonlyRecord<TName, TValue> => {
  return {
    [name]: value,
  } as ReadonlyRecord<TName, TValue>;
};

But because of this you can get the fully typed union of keys/values for the given key (and key subset):

const { contentTypesClient } = queryClients(queryClient);
const queryKey = contentTypeKeys.filterKeys.allForEnvironment(projectEnvironmentId);
const fullyTyped = contentTypesClient.getQueriesData({ queryKey });

So the core is the per-entity-query type, that holds information about all keys and values the key can represent. And this is it, mighty QueryKeyMap implementation:

export type QueryKey = readonly [ReadonlyRecord<string, unknown>];

type QueryKeyCreator = (...args: readonly any[]) => QueryKey;

type KeyName = string;

type CacheModel = ReadonlyRecord<string, any>;

type KeyGroup<TKeyShape extends QueryKey | QueryKeyCreator = never> = ReadonlyRecord<
  KeyName,
  TKeyShape
>;

type KeyGroups<TKeyShape extends QueryKey | QueryKeyCreator = never> = {
  readonly filterKeys: KeyGroup<TKeyShape>;
  readonly queryKeys: KeyGroup<TKeyShape>;
};

/**
 * Traverse through key creators and convert QueryKeyCreator fc to a final QueryKey.
 *
 * Example
 * Input TKeyGroups: { filterKeys: { filterKey1: (a,b) => [a,b] }, queryKeys: { queryKey1: (a,b,c) => [a,b,c] } }
 * Output: { filterKeys: { filterKey1: [a,b] }, queryKeys: { queryKey1: [a,b,c] } }
 */
type KeyGroupNameToKeyNameAndShape<TKeyGroups extends KeyGroups<QueryKey | QueryKeyCreator>> = {
  readonly [TKeyGroupName in keyof TKeyGroups]: TKeyGroups[TKeyGroupName] extends infer TKeyGroup
    ? {
        readonly [TKeyName in keyof TKeyGroup]: TKeyGroup[TKeyName] extends QueryKey
          ? TKeyGroup[TKeyName]
          : TKeyGroup[TKeyName] extends QueryKeyCreator
            ? ReturnType<TKeyGroup[TKeyName]>
            : never;
      }
    : never;
};

/**
 * Every query key must extend at least one filter key.
 * Also, single filter key can be extended by many query keys.
 * Using this knowledge, this mapped type simply iterates over all query key shapes and
 * filters those that extends received key shape.
 *
 * Example for TKeyShape being a queryKey
 * Input TKeyShape: [a,b,c]
 * Input TQueryKeyShapesUnion: [a,b,c] | [a,b,d]
 * Output: [a,b,c]
 *
 * Example for TKeyShape being a filterKey
 * Input TKeyShape: [a,b]
 * Input TQueryKeyShapesUnion: [a,b,c] | [a,b,d]
 * Output: [a,b,c] | [a,b,d]
 */
type GetQueryKeysForKey<
  TKeyShape extends QueryKey,
  TQueryKeyShapesUnion extends QueryKey,
> = TQueryKeyShapesUnion extends infer TQueryKey extends TKeyShape ? TQueryKey : never;

/**
 * Returns a map of all query key names to their key shapes.
 *
 * Example
 * Input TKeyGroups: { filterKeys: { filterKey1: [a,b] }, queryKeys: { queryKey1: [a,b,c] } }
 * Output: { queryKey1: [a,b,c] }
 */
type QueryKeyNamesToKeyShape<TKeyGroups extends KeyGroups<QueryKey>> =
  TKeyGroups['queryKeys'] extends infer TQueryKeys
    ? { readonly [TQueryKeyName in keyof TQueryKeys]: TQueryKeys[TQueryKeyName] }
    : never;

/**
 * Creates a map of all key names to union of query keys the key name can ever point to.
 *
 * Example
 * Input TAllKeyNamesToKeyShape: { filterKey1: [a,b], queryKey1: [a,b,c] }
 * Input TKeyGroups: { filterKeys: { filterKey1: [a,b] }, queryKeys: { queryKey1: [a,b,c] } }
 * Output: { filterKey1: [a,b,c], queryKey1: [a,b,c] }
 */
type KeyNamesToQueryKeys<
  TAllKeyNamesToKeyShape extends ReadonlyRecord<KeyName, QueryKey>,
  TKeyGroups extends KeyGroups<QueryKey>,
> = QueryKeyNamesToKeyShape<TKeyGroups> extends infer TQueryKeyNamesToKeyShape extends
  ReadonlyRecord<KeyName, QueryKey>
  ? UnionValues<TQueryKeyNamesToKeyShape> extends infer TQueryKeyShapesUnion extends QueryKey
    ? {
        readonly [TKeyName in keyof TAllKeyNamesToKeyShape]: GetQueryKeysForKey<
          TAllKeyNamesToKeyShape[TKeyName],
          TQueryKeyShapesUnion
        >;
      }
    : never
  : never;

/**
 * Flattens key groups and returns a map of all key names to their key shape.
 *
 * Example
 * Input TKeyGroups: { filterKeys: { filterKey1: [a,b] }, queryKeys: { queryKey1: [a,b,c] } }
 * Output: { filterKey1: [a,b], queryKey1: [a,b,c] }
 */
type AllKeyNamesToKeyShape<TKeyGroups extends KeyGroups<QueryKey>> = TKeyGroups['filterKeys'] &
  TKeyGroups['queryKeys'] extends infer TAllKeys
  ? { readonly [TKeyName in keyof TAllKeys]: TAllKeys[TKeyName] }
  : never;

/**
 * Accepts a map of query key names to their cache model and a map of all key names to their key shape.
 * Transfers into a tuple union of query key shapes and their cache model.
 *
 * Example
 * Input TAllKeyNamesToKeyShape: { filterKey1: [a,b], queryKey1: [a,b,c1], queryKey2: [a,b,c2] }
 * Input TQueryKeyNamesToCacheModel: { queryKey1: ABC1Model, queryKey2: ABC2Model }
 * Output: [[a,b,c1], ABC1Model] | [[a,b,c2], ABC2Model]
 */
type QueryKeyToCacheModelForAllQueryKeys<
  TQueryKeyNamesToCacheModel extends ReadonlyRecord<KeyName, CacheModel>,
  TAllKeyNamesToKeyShape extends ReadonlyRecord<string, QueryKey>,
> = keyof TQueryKeyNamesToCacheModel extends infer TKeyName
  ? TKeyName extends keyof TAllKeyNamesToKeyShape & keyof TQueryKeyNamesToCacheModel
    ? [QueryKey: TAllKeyNamesToKeyShape[TKeyName], CacheModel: TQueryKeyNamesToCacheModel[TKeyName]]
    : never // All keys from TQueryKeyNamesToCacheModel must have matching keys in TAllKeyNamesToKeyShape.
  : never;

/**
 * For the given TQueryKey returns a tuple [queryKey, cache model].
 *
 * Example for TQueryKey being a single query key
 * Input TQueryKey: [a,b,c1]
 * Input TQueryKeyToCacheModelTupleUnion: [[a,b,c1], ABC1Model] | [[a,b,c2], ABC2Model]
 * Output: [[a,b,c1], ABC1Model]
 *
 * Example for TQueryKey being a union of query keys
 * Input TQueryKey: [a,b,c1] | [a,b,c2]
 * Input TQueryKeyToCacheModelTupleUnion: [[a,b,c1], ABC1Model] | [[a,b,c2], ABC2Model]
 * Output: [[a,b,c1], ABC1Model] | [[a,b,c2], ABC2Model]
 */
type QueryKeyToCacheModelForSpecificQueryKey<
  TQueryKey extends QueryKey,
  TQueryKeyToCacheModelTupleUnion extends [QueryKey, unknown],
> = TQueryKeyToCacheModelTupleUnion extends [TQueryKey, infer TCacheModel]
  ? TQueryKeyToCacheModelTupleUnion extends [infer TKey, any]
    ? [QueryKey: TKey, CacheModel: TCacheModel]
    : never
  : never;

/**
 * Creates a map of relations between filter keys, query keys and cache models for the entity.
 * The map is then used for deriving types when working with the query cache.
 * As an example, we can get many cache records as a result of a filter query. But to keep track of what models we can receive/update,
 * we must derive the type based on the specific filter key.
 */
export type QueryKeyMap<
  TKeyGroups extends KeyGroups<QueryKey | QueryKeyCreator>,
  TQueryKeyNameToCacheModel extends ReadonlyRecord<keyof TKeyGroups['queryKeys'], CacheModel>,
> = KeyGroupNameToKeyNameAndShape<TKeyGroups> extends infer TKeyGroupNameToKeyNameAndShape extends
  KeyGroups<QueryKey>
  ? AllKeyNamesToKeyShape<TKeyGroupNameToKeyNameAndShape> extends infer TAllKeyNamesToKeyShape extends
      ReadonlyRecord<KeyName, QueryKey>
    ? KeyNamesToQueryKeys<
        TAllKeyNamesToKeyShape,
        TKeyGroupNameToKeyNameAndShape
      > extends infer TKeyNameToQueryKeys extends ReadonlyRecord<KeyName, QueryKey>
      ? {
          readonly [TKeyName in keyof TKeyNameToQueryKeys]: {
            readonly filterKey: TKeyName extends keyof TAllKeyNamesToKeyShape
              ? TAllKeyNamesToKeyShape[TKeyName]
              : never;
            readonly queryKey: TKeyNameToQueryKeys[TKeyName];
            readonly queryKeyToCacheModel: QueryKeyToCacheModelForSpecificQueryKey<
              TKeyNameToQueryKeys[TKeyName],
              QueryKeyToCacheModelForAllQueryKeys<TQueryKeyNameToCacheModel, TAllKeyNamesToKeyShape>
            >;
          };
        }
      : never
    : never
  : never;

export type EntityQueryKeyMap = ReadonlyRecord<
  KeyName,
  {
    readonly filterKey: QueryKey;
    readonly queryKey: QueryKey;
    readonly queryKeyToCacheModel: readonly [QueryKey, CacheModel];
  }
>;

export type AllKeysUnion<TEntityQueryKeyMap extends EntityQueryKeyMap> =
  TEntityQueryKeyMap[keyof TEntityQueryKeyMap]['filterKey' | 'queryKey'];

export type QueryKeysUnion<TEntityQueryKeyMap extends EntityQueryKeyMap> =
  TEntityQueryKeyMap[keyof TEntityQueryKeyMap]['queryKey'];

export type QueryKeyToCacheModelTuplesForKey<
  TAnyKey extends QueryKey,
  TKeyNameToShapes extends EntityQueryKeyMap,
> = TKeyNameToShapes[keyof TKeyNameToShapes] extends infer TEntries
  ? TEntries extends {
      readonly filterKey: TAnyKey;
      readonly queryKeyToCacheModel: infer TQueryKeyToCacheModelTuple;
    }
    ? TQueryKeyToCacheModelTuple
    : never
  : never;

export type KeyToCacheModelOrUndefinedTuple<TAnyKeyToCacheModelTuple extends [QueryKey, unknown]> =
  TAnyKeyToCacheModelTuple extends [infer TAnyKey, infer TCacheModel]
    ? [TAnyKey, TCacheModel | undefined]
    : never;

export type CacheModelForKey<
  TAnyKey extends QueryKey,
  TKeyNameToShapes extends EntityQueryKeyMap,
> = QueryKeyToCacheModelTuplesForKey<TAnyKey, TKeyNameToShapes>[1];

Feel free to implement it right into the RQ!