fix tag based invalidation by storing tags in query client

Author: KurtGokhan

.gitignore

@@ -38,4 +38,4 @@ yarn-error.log*
 next-env.d.ts
 
 # vite/vitest
-tsconfig.vitest-temp.json
+*.vitest-temp.json

README.md

@@ -32,6 +32,7 @@ It uses the builder pattern, the best pattern that works with complex Typescript
 ## TODO
 
 - Strong typed customizable tags
+- Infinite queries
 
 ## Examples
 

src/builder/QueryBuilder.ts

@@ -17,7 +17,7 @@ import {
   useSuspenseQuery,
 } from '@tanstack/react-query';
 import { operateOnTags } from '../tags/operateOnTags';
-import { QueryTagCache, QueryTagObject, QueryTagOption } from '../tags/types';
+import { QueryTagOption } from '../tags/types';
 import { FunctionType } from '../types/utils';
 import { MutationBuilder } from './MutationBuilder';
 import { MiddlewareFn, createMiddlewareFunction } from './createMiddlewareFunction';
@@ -54,7 +54,7 @@ export class QueryBuilderFrozen<T extends BuilderTypeTemplate> {
 
   getQueryFn: () => QueryFunction<T['data'], [T['vars']]> = () => {
     return ({ client, meta, queryKey, signal, pageParam }) => {
-      return this.config.queryFn({ client, meta, queryKey, signal, pageParam });
+      return this.config.queryFn({ client, meta, queryKey, signal, pageParam, originalQueryKey: queryKey });
     };
   };
 
@@ -229,9 +229,6 @@ export class QueryBuilder<T extends BuilderTypeTemplate = BuilderTypeTemplate> e
   static tagCacheId = 0;
 
   withTags(...tags: QueryTagOption<T['vars'], T['data'], T['error']>[]): this {
-    const tagCache: QueryTagCache = { tags: [] };
-    const setTags = (tags: QueryTagObject[]) => (tagCache.tags = tags);
-
     if (this.config.queryClient) {
       this.config.syncChannel?.addEventListener('message', (event) => {
         const { type, data } = event.data;
@@ -242,9 +239,7 @@ export class QueryBuilder<T extends BuilderTypeTemplate = BuilderTypeTemplate> e
       });
     }
 
-    return this.withMiddleware(createTagMiddleware<T>(tags, setTags), {
-      options: { meta: { tagCaches: { [QueryBuilder.tagCacheId++]: tagCache } } },
-    }) as unknown as this;
+    return this.withMiddleware(createTagMiddleware<T>(tags, QueryBuilder.tagCacheId++)) as unknown as this;
   }
 
   freeze(): QueryBuilderFrozen<T> {

src/builder/createMiddlewareFunction.ts

@@ -1,7 +1,6 @@
-import { QueryFunction, QueryFunctionContext } from '@tanstack/react-query';
 import { MutationBuilderConfig } from './MutationBuilder';
 import { QueryBuilderConfig } from './QueryBuilder';
-import { BuilderQueryFn, BuilderTypeTemplate } from './types';
+import type { BuilderQueryContext, BuilderQueryFn, BuilderTypeTemplate, SetAllTypes } from './types';
 
 export type MiddlewareFn<TVars, TData, TError, TOriginalTemplate extends BuilderTypeTemplate> = (
   context: MiddlewareContext<TVars>,
@@ -10,15 +9,15 @@ export type MiddlewareFn<TVars, TData, TError, TOriginalTemplate extends Builder
   config: QueryBuilderConfig<TOriginalTemplate> | MutationBuilderConfig<TOriginalTemplate>,
 ) => TData | Promise<TData>;
 
-export type MiddlewareContext<TVars> = QueryFunctionContext<[TVars]> & {
+export type MiddlewareContext<TVars> = BuilderQueryContext<TVars> & {
   vars: TVars;
 };
 
 export type MiddlewareNextFn<T extends BuilderTypeTemplate> = (
   context: Omit<MiddlewareContext<T['vars']>, 'queryKey'>,
 ) => Promise<T['data']>;
 
-const createMiddlewareContext = <TVars>(context: QueryFunctionContext<[TVars]>): MiddlewareContext<TVars> => {
+const createMiddlewareContext = <TVars>(context: BuilderQueryContext<TVars>): MiddlewareContext<TVars> => {
   return {
     ...context,
     vars: context.queryKey[0],
@@ -33,7 +32,7 @@ export const createMiddlewareFunction = <TVars, TData, TError, TOriginalTemplate
   originalFn: BuilderQueryFn<TOriginalTemplate>,
   middleware: MiddlewareFn<TVars, TData, TError, TOriginalTemplate>,
   config: QueryBuilderConfig<TOriginalTemplate> | MutationBuilderConfig<TOriginalTemplate>,
-): QueryFunction<TData, [TVars]> => {
+): BuilderQueryFn<SetAllTypes<TOriginalTemplate, TData, TError, TVars, true>> => {
   return async (context) =>
     middleware(
       createMiddlewareContext<TVars>(context),

src/builder/createTagMiddleware.ts

@@ -1,14 +1,15 @@
+import { hashKey } from '@tanstack/react-query';
 import { resolveTags } from '../tags/resolveTags';
-import { QueryTagObject, QueryTagOption } from '../tags/types';
+import { QueryTagCache, QueryTagObject, QueryTagOption } from '../tags/types';
 import { MiddlewareFn } from './createMiddlewareFunction';
 import { BuilderTypeTemplate } from './types';
 
 type CreateTagMiddleware = <T extends BuilderTypeTemplate>(
   tags: QueryTagOption[],
-  setTags: (tags: QueryTagObject[]) => void,
+  cacheId: string | number,
 ) => MiddlewareFn<T['vars'], T['data'], T['error'], T>;
 
-export const createTagMiddleware: CreateTagMiddleware = (tags, setTags) =>
+export const createTagMiddleware: CreateTagMiddleware = (tags, cacheId) =>
   async function tagMiddlware(ctx, next) {
     let resolvedTags: QueryTagObject[] = [];
 
@@ -23,6 +24,9 @@ export const createTagMiddleware: CreateTagMiddleware = (tags, setTags) =>
 
       throw error;
     } finally {
-      setTags(resolvedTags);
+      const hash = hashKey(ctx.originalQueryKey);
+      const tagCache = ((ctx.client as any).tagCache ??= {}) as QueryTagCache;
+      const tags = (tagCache[hash] ??= {});
+      tags[cacheId] = resolvedTags;
     }
   };

src/builder/types.ts

@@ -1,4 +1,4 @@
-import type { QueryFunction, UseQueryResult } from '@tanstack/react-query';
+import type { QueryFunctionContext, UseQueryResult } from '@tanstack/react-query';
 import type { RequestError } from '../http/errors';
 import type { HttpRequestOptions } from '../http/types';
 import type { Prettify } from '../types/utils';
@@ -34,7 +34,13 @@ export type BuilderQueriesResult<T extends BuilderTypeTemplate> = QueriesResultI
   queryMap: QueriesResultMapType<T>;
 };
 
-export type BuilderQueryFn<T extends BuilderTypeTemplate> = QueryFunction<T['data'], [T['vars']]>;
+export type BuilderQueryContext<TVars> = QueryFunctionContext<[TVars], never> & {
+  originalQueryKey: unknown[];
+};
+
+export type BuilderQueryFn<T extends BuilderTypeTemplate> = (
+  context: BuilderQueryContext<T['vars']>,
+) => T['data'] | Promise<T['data']>;
 
 export type SetDataType<T extends BuilderTypeTemplate, TData> = Prettify<Omit<T, 'data'> & { data: TData }>;
 

src/builder/utils.ts

@@ -36,11 +36,7 @@ export function mergeQueryOptions<
   for (const { enabled, meta, ...opt } of filtered) {
     Object.assign(opts, opt);
     opts.enabled = mergeQueryEnabled([opts.enabled, enabled]) as TOpt['enabled'];
-    opts.meta = {
-      ...opts.meta,
-      ...meta,
-      tagCaches: { ...opts.meta?.tagCaches!, ...meta?.tagCaches! },
-    };
+    opts.meta = { ...opts.meta, ...meta };
   }
 
   return opts;

src/tags/operateOnTags.ts

@@ -1,11 +1,12 @@
-import type {
-  InvalidateOptions,
-  InvalidateQueryFilters,
-  Query,
-  QueryClient,
-  QueryFilters,
+import {
+  type InvalidateOptions,
+  type InvalidateQueryFilters,
+  type Query,
+  type QueryClient,
+  type QueryFilters,
+  hashKey,
 } from '@tanstack/react-query';
-import type { QueryTag, QueryTagCache } from './types';
+import type { QueryTag } from './types';
 
 /**
  * Based on the behavior described in https://redux-toolkit.js.org/rtk-query/usage/automated-refetching#tag-invalidation-behavior
@@ -24,9 +25,11 @@ function tagMatchesTag(queryTag: QueryTag, comparedTag: QueryTag) {
   return cId === qId;
 }
 
-export function queryMatchesTag(query: Query, tag: QueryTag) {
-  const tagCaches = query.meta?.tagCaches as Record<string, QueryTagCache>;
-  const tagsInMeta = Object.values(tagCaches)?.flatMap((t) => t.tags || []);
+export function queryMatchesTag(queryClient: QueryClient, query: Query, tag: QueryTag) {
+  const hash = hashKey(query.queryKey);
+  const tagCaches = (queryClient as any)?.tagCache?.[hash] as Record<string, QueryTag[]>;
+  if (!tagCaches) return false;
+  const tagsInMeta = Object.values(tagCaches)?.flatMap((t) => t || []);
   if (tagsInMeta?.length) return tagsInMeta.some((t: QueryTag) => tagMatchesTag(t, tag));
   return false;
 }
@@ -57,7 +60,7 @@ export function operateOnTags(
     ...(operation === 'refetch' && { type: 'active' }),
     ...filters,
     predicate: (query) =>
-      tags.some((tag) => queryMatchesTag(query, tag)) && (!filters?.predicate || filters.predicate(query)),
+      tags.some((tag) => queryMatchesTag(queryClient, query, tag)) && (!filters?.predicate || filters.predicate(query)),
   };
 
   if (operation === 'refetch') return queryClient.refetchQueries(filtersObj, options);

src/tags/types.ts

@@ -81,4 +81,4 @@ export type QueryUpdateTag<TVars = unknown, TData = unknown, TErr = unknown, TTa
   QueryUpdateTagObject<TVars, TData, TErr, TTarget>
 >;
 
-export type QueryTagCache = { tags: QueryTagObject[] };
+export type QueryTagCache = Record<string | number, Record<string, QueryTagObject[]>>;

src/tags/updateTags.ts

@@ -29,7 +29,7 @@ export function updateTags({
     const tag = tags[i];
 
     const list = queryClient.getQueryCache().findAll({
-      predicate: (query) => queryMatchesTag(query, tag),
+      predicate: (query) => queryMatchesTag(queryClient, query, tag),
       type: 'all',
     });