infer data from primary tag

KurtGokhan · KurtGokhan · commit 987b3b2a4a24 · 2025-04-19T14:16:05.000+03:00
diff --git a/examples/vite/src/examples/main/example.tsx b/examples/vite/src/examples/main/example.tsx
@@ -12,20 +12,18 @@ export const builder = new HttpQueryBuilder().withClient(queryClient).withBaseUr
 export const resetMutation = builder.withPath('/reset').withMethod('post').withUpdates('*');
 
 export const articlesQuery = builder
-  .withTags('refreshable', 'articles')
+  .withTags('articles', 'refreshable')
   .withPath('/articles')
-  .withData<ArticleData[]>()
   .withSearch<{ page?: number }>()
   .withPagination({
     initialPageParam: { search: { page: 0 } },
     getNextPageParam: (prev, __, lastVars) => (!prev?.length ? null : { search: { page: (lastVars?.search?.page || 0) + 1 } }),
   });
 
 export const { useQuery: useArticle, ...articleQuery } = builder
-  .withTags('refreshable')
-  .withPath('/articles/:id')
-  .withData<ArticleData>()
+  .withTags('article', 'refreshable')
   .withTags((ctx) => ({ type: 'article', id: ctx.data.id }))
+  .withPath('/articles/:id')
   .withPreprocessor<{ id: number }>((vars) => ({ ...vars, params: { id: vars.id } }))
   .withMiddleware(async (ctx, next) => {
     const res = await next(ctx);
diff --git a/packages/tanstack-query-builder/src/builder/QueryBuilder.ts b/packages/tanstack-query-builder/src/builder/QueryBuilder.ts
@@ -1,6 +1,6 @@
 import { QueryClient } from '@tanstack/react-query';
 import { operateOnTags } from '../tags/operateOnTags';
-import type { QueryTagObject, QueryTagOption, QueryUpdateTagObject } from '../tags/types';
+import type { InferDataFromQueryTagOption, QueryTagObject, QueryTagOption, QueryUpdateTagObject } from '../tags/types';
 import { TODO } from '../type-utils';
 import { QueryBuilderFrozen } from './QueryBuilderFrozen';
 import { type MiddlewareFn, createMiddlewareFunction } from './createMiddlewareFunction';
@@ -112,15 +112,18 @@ export class QueryBuilder<
    * Adds a tag to the query.
    * The tag is used to invalidate or update the query when the tag is updated.
    */
-  withTags(...tags: QueryTagOption<TVars, TData, TError, QueryTagObject<TTags>>[]): this {
-    return this.withMiddleware(createTagMiddleware(tags.flat(), QueryBuilder.tagCacheId++)) as unknown as this;
+  withTags<const TTagOpt extends QueryTagOption<TVars, TData, TError, QueryTagObject<TTags>, TTags>>(
+    primaryTag: TTagOpt,
+    ...otherTags: QueryTagOption<TVars, TData, TError, QueryTagObject<TTags>, TTags>[]
+  ): unknown extends TData ? QueryBuilder<TVars, InferDataFromQueryTagOption<TTagOpt, TTags>, TError, TKey, TTags, TFlags> : this {
+    return this.withMiddleware(createTagMiddleware([primaryTag, ...otherTags].flat(), QueryBuilder.tagCacheId++)) as unknown as any;
   }
 
   /**
    * Adds a declarative update to the mutation.
    * This is used to invalidate or update the queries that were marked with {@link withTags}.
    */
-  withUpdates(...tags: QueryTagOption<TVars, TData, TError, QueryUpdateTagObject<TVars, TData, TError, TTags>>[]): this {
+  withUpdates(...tags: QueryTagOption<TVars, TData, TError, QueryUpdateTagObject<TVars, TData, TError, TTags>, TTags>[]): this {
     return this.withMiddleware(createUpdateMiddleware<TVars, TData, TError, TKey, TTags>(tags)) as unknown as this;
   }
 
diff --git a/packages/tanstack-query-builder/src/builder/QueryBuilderClient.ts b/packages/tanstack-query-builder/src/builder/QueryBuilderClient.ts
@@ -40,7 +40,7 @@ export class QueryBuilderClient<
   TError,
   TKey extends unknown[],
   TTags extends Record<string, unknown>,
-  TFilters = QueryFilters<TData, TError, TData, TKey>,
+  TFilters = QueryFilters<TKey>,
 > {
   private declare _options: BuilderConfig<TVars, TData, TError, TKey>['options'];
   private declare _pgOptions: BuilderConfig<TVars, TData, TError, TKey>['paginationOptions'];
diff --git a/packages/tanstack-query-builder/src/builder/createTagMiddleware.ts b/packages/tanstack-query-builder/src/builder/createTagMiddleware.ts
@@ -4,7 +4,7 @@ import type { QueryTagCache, QueryTagObject, QueryTagOption } from '../tags/type
 import type { MiddlewareFn } from './createMiddlewareFunction';
 
 type CreateTagMiddleware = <TVars, TData, TError, TKey extends unknown[]>(
-  tags: QueryTagOption<any, any, any>[],
+  tags: QueryTagOption<any, any, any, any>[],
   cacheId: string | number,
 ) => MiddlewareFn<TVars, TData, TError, TKey>;
 
diff --git a/packages/tanstack-query-builder/src/builder/createUpdateMiddleware.ts b/packages/tanstack-query-builder/src/builder/createUpdateMiddleware.ts
@@ -5,7 +5,7 @@ import { type UpdateTagsUndoer, updateTags } from '../tags/updateTags';
 import type { MiddlewareFn } from './createMiddlewareFunction';
 
 type CreateUpdateMiddleware = <TVars, TData, TError, TKey extends unknown[], TTags extends Record<string, unknown>>(
-  tags: QueryTagOption<TVars, TData, TError, QueryUpdateTagObject<TVars, TData, TError, TTags>>[],
+  tags: QueryTagOption<TVars, TData, TError, QueryUpdateTagObject<TVars, TData, TError, TTags>, TTags>[],
 ) => MiddlewareFn<TVars, TData, TError, TKey>;
 
 export const createUpdateMiddleware: CreateUpdateMiddleware = (tags) =>
diff --git a/packages/tanstack-query-builder/src/http/HttpQueryBuilder.ts b/packages/tanstack-query-builder/src/http/HttpQueryBuilder.ts
@@ -3,6 +3,7 @@ import { QueryBuilder } from '../builder/QueryBuilder';
 import type { BuilderPaginationOptions } from '../builder/options';
 import type { BuilderConfig, BuilderFlags } from '../builder/types';
 import type { RequestError } from '../http/errors';
+import { InferDataFromQueryTagOption, QueryTagObject, QueryTagOption } from '../tags/types';
 import type { WithOptional } from '../type-utils';
 import type { HttpBaseHeaders, HttpBaseParams, HttpBaseSearch, HttpBuilderVars } from './builder-types';
 import { createHttpMergeVarsFn, createHttpQueryFn, createHttpQueryKeySanitizer } from './builder-utils';
@@ -128,6 +129,21 @@ export class HttpQueryBuilder<
     paginationOptions: BuilderPaginationOptions<HttpBuilderVars<TParam, TSearch, TBody, THeader, TMeta>, TData, TError, TKey>,
   ) => HttpQueryBuilder<TParam, TSearch, TBody, THeader, TMeta, TData, TError, TTags, TFlags | 'withPagination', TKey>;
 
+  declare withTags: <
+    const TTagOpt extends QueryTagOption<
+      HttpBuilderVars<TParam, TSearch, TBody, THeader, TMeta>,
+      TData,
+      TError,
+      QueryTagObject<TTags>,
+      TTags
+    >,
+  >(
+    primaryTag: TTagOpt,
+    ...otherTags: QueryTagOption<HttpBuilderVars<TParam, TSearch, TBody, THeader, TMeta>, TData, TError, QueryTagObject<TTags>, TTags>[]
+  ) => unknown extends TData
+    ? HttpQueryBuilder<TParam, TSearch, TBody, THeader, TMeta, InferDataFromQueryTagOption<TTagOpt, TTags>, TError, TTags, TFlags, TKey>
+    : this;
+
   withTagTypes<TTag extends string, T = unknown>(): HttpQueryBuilder<
     TParam,
     TSearch,
diff --git a/packages/tanstack-query-builder/src/tags/resolveTags.ts b/packages/tanstack-query-builder/src/tags/resolveTags.ts
@@ -2,7 +2,7 @@ import type { QueryClient } from '@tanstack/react-query';
 import type { QueryTag, QueryTagContext, QueryTagObject, QueryTagOption } from './types';
 
 export type ResolveTagsParams<TVars = void, TData = unknown, TErr = unknown, TTag extends QueryTagObject = QueryTagObject> = {
-  tags?: QueryTagOption<TVars, TData, TErr, TTag> | QueryTagOption<TVars, TData, TErr, TTag>[] | null;
+  tags?: QueryTagOption<TVars, TData, TErr, TTag, any> | QueryTagOption<TVars, TData, TErr, TTag, any>[] | null;
   vars: TVars;
   data?: TData;
   error?: unknown;
@@ -21,7 +21,7 @@ export function resolveTags<TVars = void, TTag extends QueryTagObject = QueryTag
 }: ResolveTagsParams<TVars, any, any, TTag> & { client: QueryClient }): TTag[] {
   const resolvedTags =
     typeof tags === 'function'
-      ? resolveTags({ tags: tags(ctx as QueryTagContext<TVars>), ...ctx })
+      ? resolveTags({ tags: tags(ctx as QueryTagContext<TVars, any>), ...ctx })
       : Array.isArray(tags)
         ? tags.flatMap((tag) => resolveTags({ tags: tag, ...ctx }))
         : // If there is an error, tags passed as array will be ignored. User should pass tags as function to handle error case if needed.
diff --git a/packages/tanstack-query-builder/src/tags/types.ts b/packages/tanstack-query-builder/src/tags/types.ts
@@ -20,15 +20,29 @@ export type QueryTagContext<TVars = void, TData = unknown, TErr = unknown> = {
   error?: TErr;
 };
 
-export type QueryTagCallback<TVars = void, TData = unknown, TErr = unknown, TTag = QueryTag> = (
-  ctx: QueryTagContext<TVars, TData, TErr>,
-) => TTag | readonly TTag[];
+export type QueryTagCallback<
+  TVars = void,
+  TData = unknown,
+  TErr = unknown,
+  TTag extends QueryTagObject<any> = QueryTagObject<any>,
+  TTags extends TagMap = TagMap,
+> = (ctx: QueryTagContext<TVars, unknown extends TData ? InferDataFromQueryTagOption<TTag, TTags> : TData, TErr>) => TTag | readonly TTag[];
 
 export type QueryTagStaticOption<TTag extends QueryTagObject<any> = QueryTagObject<any>> = '*' | QueryTag<TTag> | readonly QueryTag<TTag>[];
 
-export type QueryTagOption<TVars = unknown, TData = unknown, TErr = unknown, TTag extends QueryTagObject<any> = QueryTagObject<any>> =
-  | QueryTagStaticOption<TTag>
-  | QueryTagCallback<TVars, TData, TErr, TTag>;
+export type QueryTagOption<
+  TVars = unknown,
+  TData = unknown,
+  TErr = unknown,
+  TTag extends QueryTagObject<any> = QueryTagObject<any>,
+  TTags extends TagMap = TagMap,
+> = QueryTagStaticOption<TTag> | QueryTagCallback<TVars, TData, TErr, TTag, TTags>;
+
+export type InferDataFromQueryTagOption<TTag extends QueryTagOption<any, any, any, any>, TTags extends TagMap> = TTag extends string
+  ? TTags[TTag]
+  : TTag extends QueryTagCallback<any, infer TData, any, any>
+    ? TData
+    : unknown;
 
 export type QueryUpdaterFn<TVars = unknown, TData = unknown, TErr = unknown, TTarget = unknown> = (
   ctx: QueryTagContext<TVars, TData, TErr>,
